diff --git a/benchmark.py b/benchmark.py index 7b19c41..5dbb42e 100644 --- a/benchmark.py +++ b/benchmark.py @@ -28,7 +28,7 @@ def benchmark_query(query: SelectQuery, filename: str = None) -> Tuple[float, fl return result["Planning Time"], result["Execution Time"] -def statistics(query: SelectQuery, filename: str, repeats: int = 500): +def statistics(query: SelectQuery, filename: str, repeats: int = 500) -> None: ts = shutil.get_terminal_size((80, 20)) print(filename.center(ts.columns, "-")) planning_times = [] diff --git a/data.py b/data.py index 0f96a24..337e68e 100644 --- a/data.py +++ b/data.py @@ -1,3 +1,6 @@ +from dataclasses import dataclass +from typing import Optional + colors_c1 = { "Travis": "#7592a4", "Marisha": "#bd6b1e", @@ -72,24 +75,32 @@ single_speaker = { assert set(single_speaker["Handbooker Helper"].keys()) == set(range(1, 44 + 1)) -series_data = [ - { - "name": "Campaign 1", - "playlist_id": "PL1tiwbzkOjQz7D0l_eLJGAISVtcL7oRu_", - }, - { - "name": "Campaign 2", - "playlist_id": "PL1tiwbzkOjQxD0jjAE7PsWoaCrs0EkBH2" - }, - { - "name": "Handbooker Helper", - "playlist_id": "PL1tiwbzkOjQyr6-gqJ8r29j_rJkR49uDN", - "single_speaker": True - }, - { - "name": "Mini Primetime", - "playlist_id": "PL1tiwbzkOjQz9kKDaPRPrX2E7RPTaxEZd", - "initial_speaker": "Will" - } +@dataclass +class SeriesData: + name: str + playlist_id: str + single_speaker: bool = False + initial_speaker: Optional[str] = None + + +series_data = [ + SeriesData( + name="Campaign 1", + playlist_id="PL1tiwbzkOjQz7D0l_eLJGAISVtcL7oRu_" + ), + SeriesData( + name="Campaign 2", + playlist_id="PL1tiwbzkOjQxD0jjAE7PsWoaCrs0EkBH2" + ), + SeriesData( + name="Handbooker Helper", + playlist_id="PL1tiwbzkOjQyr6-gqJ8r29j_rJkR49uDN", + single_speaker=True + ), + SeriesData( + name="Mini Primetime", + playlist_id="PL1tiwbzkOjQz9kKDaPRPrX2E7RPTaxEZd", + initial_speaker="Will" + ), ] diff --git a/fetch.py b/fetch.py index 67c000f..d3b5923 100644 --- a/fetch.py +++ b/fetch.py @@ -13,11 +13,11 @@ from models import Episode, Series, Line, Phrase from utils import srtdir, pretty_title -def main(): +def main() -> None: os.nice(15) for series in series_data: - name = series["name"] - playlist_id = series["playlist_id"] + name = series.name + playlist_id = series.playlist_id is_campaign = "Campaign" in name try: s = Series.select().where(Series.title == name).get() @@ -26,7 +26,7 @@ def main(): s.title = name s.is_campaign = is_campaign - s.single_speaker = "single_speaker" in series and series["single_speaker"] + s.single_speaker = series.single_speaker s.save() ydl_opts = { 'extract_flat': True @@ -37,7 +37,7 @@ def main(): print(v["url"] for v in videos) - ydl_opts = { + ydl_opts_download = { "writesubtitles": True, "subtitleslangs": ["en", "en-US"], "skip_download": True, @@ -58,6 +58,8 @@ def main(): if s.is_campaign: try: match = regex.search(video["title"]) + if not match: + raise ValueError("No episode number found in title") e.episode_number = int(match.group(1)) except AttributeError: if s.title == "Campaign 1": # one-shots at the end of campaign 1 @@ -71,8 +73,8 @@ def main(): print(e.series.id, e.episode_number, e.pretty_title) vttfile = srtdir / str(e.id) - ydl_opts["outtmpl"] = str(vttfile) - with youtube_dl.YoutubeDL(ydl_opts) as ydl: + ydl_opts_download["outtmpl"] = str(vttfile) + with youtube_dl.YoutubeDL(ydl_opts_download) as ydl: ydl.download([f'https://www.youtube.com/watch?v={e.youtube_id}']) if vttfile.with_suffix(".en-US.vtt").exists(): # few videos have en-US as language code instead of en diff --git a/import.py b/import.py index 34b4a99..f916a9a 100644 --- a/import.py +++ b/import.py @@ -1,7 +1,7 @@ import os import re from html import unescape -from typing import List, Optional +from typing import List, Optional, Set from alive_progress import alive_bar from peewee import fn, chunked @@ -28,7 +28,7 @@ def add_to_text(text: str, add: str) -> str: return add -def insert_subtitle(text: str, person: str, subline: Subtitle, episode: Episode, order: int, +def insert_subtitle(text: str, person: Optional[Person], subline: Subtitle, episode: Episode, order: int, isnote: bool = False, ismeta: bool = False) -> Line: dbline = Line() if not text: @@ -46,9 +46,9 @@ def insert_subtitle(text: str, person: str, subline: Subtitle, episode: Episode, return dbline -def main(): +def main() -> None: os.nice(15) - all_people = set() + all_people: Set[str] = set() for series in Series.select().order_by(Series.id): for episode in Episode.select().where( (Episode.text_imported == False) & (Episode.series == series) & (Episode.downloaded) diff --git a/mypy.ini b/mypy.ini new file mode 100644 index 0000000..7fb1515 --- /dev/null +++ b/mypy.ini @@ -0,0 +1,8 @@ +[mypy] +warn_return_any = True +warn_unused_configs = True +disallow_untyped_defs = True +disallow_incomplete_defs = True +ignore_missing_imports = True +check_untyped_defs=True +warn_redundant_casts=True diff --git a/server.py b/server.py index 4417824..415568c 100644 --- a/server.py +++ b/server.py @@ -1,4 +1,3 @@ -import logging from typing import List from flask import request, jsonify, Response @@ -10,12 +9,13 @@ from psycopg2._psycopg import cursor from app import app from models import * + # logger = logging.getLogger('peewee') # logger.addHandler(logging.StreamHandler()) # logger.setLevel(logging.DEBUG) -def add_cors(response): +def add_cors(response: Response) -> Response: header = response.headers header['Access-Control-Allow-Origin'] = '*' return response