import logging
import os
from collections.abc import Sequence
from concurrent.futures import Executor, ThreadPoolExecutor
from datetime import datetime
from pathlib import Path
import re
from typing import Optional, Self
import dateutil.parser
import tenacity
from requests import Response
from rich.console import Group
from rich.live import Live
from rich.panel import Panel
from rich.progress import (
BarColumn,
DownloadColumn,
MofNCompleteColumn,
Progress,
TaskID,
TaskProgressColumn,
TextColumn,
TimeElapsedColumn,
TimeRemainingColumn,
TransferSpeedColumn,
)
from rich.style import Style, StyleType
from thu_learn_downloader.client.client import Client
from thu_learn_downloader.client.course import Course
from thu_learn_downloader.client.document import Document, DocumentClass
from thu_learn_downloader.client.homework import Homework
from thu_learn_downloader.client.semester import Semester
from . import description, filename, style
from .selector import Selector
class Downloader:
prefix: Path
selector: Selector
executor: Executor
live: Live
semesters_task_id: TaskID
courses_task_id: TaskID
documents_task_id: TaskID
def __init__(
self,
prefix: Optional[Path] = None,
selector: Optional[Selector] = None,
jobs: int = 8,
) -> None:
self.prefix = prefix or Path.home() / "thu-learn"
self.selector = selector or Selector()
self.executor = ThreadPoolExecutor(max_workers=jobs)
self.progress_prepare = Progress(
TextColumn("{task.description}", style="bold bright_blue"),
BarColumn(),
MofNCompleteColumn(),
TimeElapsedColumn(),
)
self.progress_download = Progress(
TextColumn("{task.description}", style="bold"),
BarColumn(),
DownloadColumn(),
TaskProgressColumn(),
TimeElapsedColumn(),
TimeRemainingColumn(),
TransferSpeedColumn(),
)
self.semesters_task_id = self.progress_prepare.add_task(description="Semesters")
self.courses_task_id = self.progress_prepare.add_task(description="Courses")
self.documents_task_id = self.progress_prepare.add_task(description="Documents")
self.live = Live(
Group(
Panel(self.progress_download, height=jobs + 2),
Panel(self.progress_prepare),
)
)
def __enter__(self) -> Self:
self.live.__enter__()
self.executor.__enter__()
return self
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
self.executor.__exit__(exc_type, exc_val, exc_tb)
self.live.__exit__(exc_type, exc_val, exc_tb)
@tenacity.retry(
stop=tenacity.stop_after_attempt(max_attempt_number=4),
before_sleep=tenacity.before_sleep_log(logging.getLogger(), logging.DEBUG),
)
def download(
self,
client: Client,
url: str,
output: Path,
*,
task_id: TaskID,
) -> None:
self.progress_download.update(task_id=task_id, visible=True)
response: Response = client.get(url=url, stream=True)
os.makedirs(name=output.parent, exist_ok=True)
with output.open(mode="wb") as file:
for chunk in response.iter_content(chunk_size=8192):
bytes_written: int = file.write(chunk)
self.progress_download.advance(task_id=task_id, advance=bytes_written)
self.progress_download.update(task_id=task_id, visible=False)
def do_sync_file(
self,
client: Client,
url: str,
output: Path,
description: str,
*,
remote_size: Optional[int] = None,
remote_time: Optional[datetime] = None,
style: Optional[StyleType] = None,
) -> None:
if remote_size is None or remote_time is None:
response: Response = client.get(url=url, stream=True)
if remote_size is None:
try:
remote_size = int(response.headers["Content-Length"])
except Exception:
remote_size = None
if remote_time is None:
try:
remote_time = dateutil.parser.parse(response.headers["Date"])
except Exception:
remote_time = None
if (
remote_size is not None
and output.exists()
and output.stat().st_size == remote_size
):
style = style or Style()
if not isinstance(style, Style):
style = Style.parse(style)
self.live.console.log(
"[reverse] SKIPPED [/]", description, style=style + Style(dim=True)
)
return
task_id: TaskID = self.progress_download.add_task(
description=description, total=remote_size
)
self.download(client=client, url=url, output=output, task_id=task_id)
if remote_time:
os.utime(
path=output, times=(remote_time.timestamp(), remote_time.timestamp())
)
self.live.console.log("[reverse] SUCCESS [/]", description, style=style)
def sync_file(
self,
client: Client,
url: str,
output: Path,
description: str,
*,
remote_size: Optional[int] = None,
remote_time: Optional[datetime] = None,
style: Optional[StyleType] = None,
) -> None:
self.executor.submit(
self.do_sync_file,
client=client,
url=url,
output=output,
description=description,
remote_size=remote_size,
remote_time=remote_time,
style=style,
)
def sync_semesters(self, semesters: Sequence[Semester]) -> None:
if self.selector.semesters:
semesters = [
semester
for semester in semesters
if semester.id in self.selector.semesters
]
self.progress_prepare.reset(task_id=self.semesters_task_id)
for semester in self.progress_prepare.track(
sequence=semesters,
total=len(semesters),
task_id=self.semesters_task_id,
description="Semesters",
):
self.sync_semester(semester=semester)
def sync_semester(self, semester: Semester) -> None:
self.sync_courses(semester=semester, courses=semester.courses)
def sync_courses(self, semester: Semester, courses: Sequence[Course]) -> None:
if self.selector.courses:
courses = [
course for course in courses if course.id in self.selector.courses
]
self.progress_prepare.reset(task_id=self.courses_task_id)
for course in self.progress_prepare.track(
sequence=courses,
total=len(courses),
task_id=self.courses_task_id,
description="Courses",
):
self.sync_course(semester=semester, course=course)
def sync_course(self, semester: Semester, course: Course) -> None:
if self.selector.document:
self.sync_documents(
semester=semester,
course=course,
document_classes=course.document_classes,
documents=course.documents,
)
if self.selector.homework:
self.sync_homeworks(
semester=semester, course=course, homeworks=course.homeworks
)
def sync_documents(
self,
semester: Semester,
course: Course,
document_classes: Sequence[DocumentClass],
documents: Sequence[Document],
) -> None:
document_class_map: dict[str, DocumentClass] = {
document_class.id: document_class for document_class in document_classes
}
self.progress_prepare.reset(task_id=self.documents_task_id)
for index, document in enumerate(
self.progress_prepare.track(
sequence=documents,
total=len(documents),
task_id=self.documents_task_id,
description="Documents",
),
start=1,
):
self.sync_document(
semester=semester,
course=course,
document_class=document_class_map[document.class_id],
document=document,
index=index,
)
def sync_document(
self,
semester: Semester,
course: Course,
document_class: DocumentClass,
document: Document,
index: int,
) -> None:
self.sync_file(
client=document.client,
url=document.download_url,
output=filename.document(
prefix=self.prefix,
semester=semester,
course=course,
document_class=document_class,
document=document,
index=index,
),
description=description.document(
semester=semester,
course=course,
document_class=document_class,
document=document,
index=index,
),
remote_size=document.size,
remote_time=document.upload_time,
style=style.DOCUMENT,
)
def sync_homeworks(
self, semester: Semester, course: Course, homeworks: Sequence[Homework]
) -> None:
self.progress_prepare.reset(task_id=self.documents_task_id)
for homework in self.progress_prepare.track(
sequence=homeworks,
total=len(homeworks),
task_id=self.documents_task_id,
description="Homeworks",
):
self.sync_homework(semester=semester, course=course, homework=homework)
def sync_homework(
self, semester: Semester, course: Course, homework: Homework
) -> None:
readme_path: Path = filename.homework(
prefix=self.prefix, semester=semester, course=course, homework=homework
)
readme_path = Path(re.sub(r'[<>:"\\|?*\x00-\x1F]', '_', str(readme_path)))
os.makedirs(readme_path.parent, exist_ok=True)
# readme_path.write_text(homework.markdown)
for attachment in homework.attachments:
self.sync_file(
client=homework.client,
url=attachment.download_url,
output=filename.attachment(
prefix=self.prefix,
semester=semester,
course=course,
homework=homework,
attachment=attachment,
),
description=description.attachment(
semester=semester,
course=course,
homework=homework,
attachment=attachment,
),
style=style.HOMEWORK,
)
