Files
MediaManager/media_manager/indexer/schemas.py
Marcel Hellwig 96b84d45db Adding some more new lints (#393)
Enable `UP` and `TRY` lint
2026-02-01 18:04:15 +01:00

94 lines
2.9 KiB
Python

import re
import typing
from uuid import UUID, uuid4
import pydantic
from pydantic import BaseModel, ConfigDict, computed_field
from media_manager.torrent.models import Quality
IndexerQueryResultId = typing.NewType("IndexerQueryResultId", UUID)
class IndexerQueryResult(BaseModel):
model_config = ConfigDict(from_attributes=True)
id: IndexerQueryResultId = pydantic.Field(
default_factory=lambda: IndexerQueryResultId(uuid4())
)
title: str
download_url: str = pydantic.Field(
exclude=True,
description="This can be a magnet link or URL to the .torrent file",
)
seeders: int
flags: list[str]
size: int
usenet: bool
age: int
score: int = 0
indexer: str | None
@computed_field
@property
def quality(self) -> Quality:
high_quality_pattern = r"\b(4k)\b"
medium_quality_pattern = r"\b(1080p)\b"
low_quality_pattern = r"\b(720p)\b"
very_low_quality_pattern = r"\b(480p|360p)\b"
if re.search(high_quality_pattern, self.title, re.IGNORECASE):
return Quality.uhd
if re.search(medium_quality_pattern, self.title, re.IGNORECASE):
return Quality.fullhd
if re.search(low_quality_pattern, self.title, re.IGNORECASE):
return Quality.hd
if re.search(very_low_quality_pattern, self.title, re.IGNORECASE):
return Quality.sd
return Quality.unknown
@computed_field
@property
def season(self) -> list[int]:
pattern = r"\bS(\d+)\b"
matches = re.findall(pattern, self.title, re.IGNORECASE)
if matches.__len__() == 2:
result = list(range(int(matches[0]), int(matches[1]) + 1))
elif matches.__len__() == 1:
result = [int(matches[0])]
else:
result = []
return result
def __gt__(self, other: "IndexerQueryResult") -> bool:
if self.quality.value != other.quality.value:
return self.quality.value < other.quality.value
if self.score != other.score:
return self.score > other.score
if self.usenet != other.usenet:
return self.usenet
if self.usenet and other.usenet:
return self.age > other.age
if not self.usenet and not other.usenet:
return self.seeders > other.seeders
return self.size < other.size
def __lt__(self, other: "IndexerQueryResult") -> bool:
if self.quality.value != other.quality.value:
return self.quality.value > other.quality.value
if self.score != other.score:
return self.score < other.score
if self.usenet != other.usenet:
return not self.usenet
if self.usenet and other.usenet:
return self.age < other.age
if not self.usenet and not other.usenet:
return self.seeders < other.seeders
return self.size > other.size