From 4c6ba6bd4a1826e5526dd9029e8048309f61b82b Mon Sep 17 00:00:00 2001
From: cinnaboot <cinnaboot@protonmail.com>
Date: Wed, 8 Apr 2020 14:25:58 -0400
Subject: [PATCH] use guessit library for title parsing

---
 video_metadata/views/tmdb.py | 79 +++++++++++++++++-------------------
 1 file changed, 38 insertions(+), 41 deletions(-)

diff --git a/video_metadata/views/tmdb.py b/video_metadata/views/tmdb.py
index 57fdb43..aabf711 100644
--- a/video_metadata/views/tmdb.py
+++ b/video_metadata/views/tmdb.py
@@ -16,6 +16,7 @@ import urllib.request
 
 from flask import Blueprint, current_app, flash, Flask, jsonify, make_response, request, \
     redirect, render_template, Response, url_for
+from guessit import guessit
 
 from .. import util
 from .. import db
@@ -53,13 +54,17 @@ class media_type(IntEnum):
 @dataclass
 class query_data:
     errors: List[str] = field(default_factory=list)
-    full_url: str = ''
-    movie_year: int = 0
-    parsed_query_str: str = ''
-    query_params: str = ''
-    media_type: media_type = media_type.UNKNOWN
-    raw_query_str: str = ''
-    rel_path: str = ''
+    is_valid:bool = True
+    raw_query_str:str = ''
+    parsed_query_str:str = ''
+    title:str = ''
+    media_type:media_type = media_type.UNKNOWN
+    rel_path:str = ''
+    query_params:str = ''
+    full_url:str = ''
+    movie_year:int = 0
+    tv_season:str = ''
+    tv_episode:str = ''
 
 @dataclass
 class response_data:
@@ -168,13 +173,14 @@ def tv_show_test():
 @bp.route('/manual_query', methods=['GET', 'POST'])
 def manual_query():
     query = query_data()
+    view_results = []
+    errors = []
 
     if 'query' in request.form:
         query.raw_query_str = request.form['query']
-
-    q_res = doQueryAlgorithm(query)
-    errors = q_res['errors']
-    view_results = q_res['results']
+        q_res = doQueryAlgorithm(query)
+        errors = q_res['errors']
+        view_results = q_res['results']
 
     if errors:
         view_results.clear()
@@ -338,10 +344,9 @@ def cache_image(image_name:str = ''):
 def doQueryAlgorithm(query:query_data) -> dict:
     errors, view_results = [], []
     response = response_data()
+    query = guessMedia(query)
 
-    # TODO: if query_obj.media_type is still 'UNKNOWN' from parseQueryStr(), make queries to
-    #   both tmdb endpoints (movie and tv)
-    if query.raw_query_str and parseQueryStr(query):
+    if query.is_valid:
         if isQueryCached(query.parsed_query_str, query.movie_year):
             view_results = makeLocalQuery(query)
             current_app.logger.debug(f'view_results length: {len(view_results)}')
@@ -370,35 +375,25 @@ def doQueryAlgorithm(query:query_data) -> dict:
     return {'errors': errors, 'results': view_results}
 
 
-# TODO: need a way to correctly parse titles like: "2001: A Space Odyssey (1968)"
-#   with legitimate dates in the title
-def parseQueryStr(query_obj):
-    s = query_obj.raw_query_str.replace('.', ' ').strip().lower()
-    s = re.sub('[(){}\[\]]', '', s).strip() # NOTE: remove some common characters
-    current_year = date.today().year
+def guessMedia(query:query_data) -> query_data:
+    guess = guessit(query.raw_query_str)
+    query.title = guess['title']
+    query.parsed_query_str = guess['title'].lower()
 
-    # NOTE: check for a string that looks like a tv show
-    m = re.search(r's[0-9]{1,2}|season ?[0-9]{1,2}', s, re.IGNORECASE)
-
-    if m:
-        query_obj.parsed_query_str = s[0:m.start()].strip()
-        query_obj.media_type = media_type.TV_SHOW
-        return True
-
-    # NOTE: find all 4 digit numbers to match the year of queries with 4 digit numbers in the title
-    #   ex) "2036.Origin.Unknown.2018"
-    regex = re.compile('[0-9]{4}')
-
-    for m in regex.finditer(s):
-        i = int(m[0])
-        if i > 1878 and i <= current_year:
-            query_obj.parsed_query_str = s[0:m.start()].strip() # NOTE: slice query string to year match
-            query_obj.movie_year = i
-            query_obj.media_type = media_type.MOVIE
-            return True
+    try:
+        if guess['type'] == 'movie':
+            query.media_type = media_type.MOVIE
+            query.movie_year = guess['year']
+        else: # NOTE: 'type' has to be 'episode' here. guessit has no 'unknown' equivalent
+            query.media_type = media_type.TV_SHOW
+            query.tv_season = guess['season']
+            query.tv_episode = guess['episode']
+    except KeyError as e:
+        query.is_valid = False
+        current_app.logger.error(f'invalid key: {e}, for query: {query.raw_query_str}')
 
-    current_app.logger.warning(f'unable to guess media type for "{query_obj.raw_query_str}"')
-    return False
+    current_app.logger.debug(f'query: {util.dumpObject(query)}')
+    return query
 
 
 def isQueryCached(query_str:str, year) -> bool:
@@ -470,6 +465,8 @@ def buildViewResult(db_row):
 
 
 # TODO: need to check response headers for rate limiting requests
+# TODO: if query_obj.media_type is still 'UNKNOWN' from parseQueryStr(), make queries to
+#   both tmdb endpoints (movie and tv)
 def makeTMDBQuery(query_obj, response_obj):
     if query_obj.media_type == media_type.MOVIE:
         search_part = TMDB_SEARCH_MOVIE