Add watching status and groups for shows and movies

tools/content/media.py

@@ -6,22 +6,39 @@ MEDIA_KINDS = {
     "show": {"key": "shows", "label": "TV show", "omdb_type": "series"},
     "movie": {"key": "movies", "label": "movie", "omdb_type": "movie"},
 }
+STATUSES = ["watched", "watching"]
 
 
-def check_rating(rating: int) -> int:
-    if not 0 <= rating <= 10:
+def check_rating(rating: int | None) -> int | None:
+    if rating is not None and not 0 <= rating <= 10:
         raise ValueError("Rating must be between 0 and 10")
     return rating
 
 
+def check_status(status: str) -> str:
+    if status not in STATUSES:
+        raise ValueError(f"Status must be one of {', '.join(STATUSES)}")
+    return status
+
+
 def load_items(kind: str) -> tuple[dict, list[dict]]:
     data = jsonfile.load(ABOUT_FILE)
     return data, data.setdefault(MEDIA_KINDS[kind]["key"], [])
 
 
+def status_of(item: dict) -> str:
+    return item.get("status", "watched")
+
+
 def list_items(kind: str) -> list[dict]:
     _, items = load_items(kind)
-    return sorted(items, key=lambda item: -item.get("rating", 0))
+    rating = lambda item: item["rating"] if item.get("rating") is not None else -1
+    return sorted(items, key=lambda item: (status_of(item) != "watching", -rating(item)))
+
+
+def collection_names(kind: str) -> list[str]:
+    _, items = load_items(kind)
+    return sorted({item["collection"] for item in items if item.get("collection")})
 
 
 def find_index(items: list[dict], text: str) -> int | None:
@@ -34,27 +51,53 @@ def find_index(items: list[dict], text: str) -> int | None:
     return None
 
 
-def look_up_name(kind: str, identifier: str) -> str:
+def look_up_name(identifier: str) -> str:
     details = imdb.lookup(identifier)
     if not details:
         raise ValueError("Could not look up the title on OMDb. Enter a name, or set OMDB_API_KEY in .env")
     return details.get("Title") or identifier
 
 
-def add_item(kind: str, text: str, rating: int, name: str | None = None) -> dict:
+def with_optional(item: dict, key: str, value) -> dict:
+    updated = {name: field for name, field in item.items() if name != key}
+    if value is not None and value != "":
+        updated[key] = value
+    return updated
+
+
+def build_item(name: str, url: str, rating: int | None, status: str, collection: str | None) -> dict:
+    item = {"name": name, "url": url}
+    item = with_optional(item, "rating", rating)
+    if status != "watched":
+        item["status"] = status
+    return with_optional(item, "collection", collection)
+
+
+def add_item(kind: str, text: str, rating: int | None, name: str | None = None, status: str = "watched", collection: str | None = None) -> dict:
     identifier = imdb.imdb_id(text)
     if not identifier:
         raise ValueError(f"{text} is not an IMDb link or ID")
     check_rating(rating)
+    check_status(status)
     data, items = load_items(kind)
     if find_index(items, identifier) is not None:
         raise ValueError(f"{identifier} is already listed")
-    item = {"name": name or look_up_name(kind, identifier), "rating": rating, "url": imdb.title_url(identifier)}
+    item = build_item(name or look_up_name(identifier), imdb.title_url(identifier), rating, status, collection)
     items.append(item)
     jsonfile.save(ABOUT_FILE, data)
     return item
 
 
+def update_item(kind: str, text: str, change) -> dict:
+    data, items = load_items(kind)
+    index = find_index(items, text)
+    if index is None:
+        raise ValueError(f"{text} is not listed")
+    items[index] = change(items[index])
+    jsonfile.save(ABOUT_FILE, data)
+    return items[index]
+
+
 def remove_item(kind: str, text: str) -> dict:
     data, items = load_items(kind)
     index = find_index(items, text)
@@ -65,12 +108,15 @@ def remove_item(kind: str, text: str) -> dict:
     return removed
 
 
-def rate_item(kind: str, text: str, rating: int) -> dict:
+def rate_item(kind: str, text: str, rating: int | None) -> dict:
     check_rating(rating)
-    data, items = load_items(kind)
-    index = find_index(items, text)
-    if index is None:
-        raise ValueError(f"{text} is not listed")
-    items[index]["rating"] = rating
-    jsonfile.save(ABOUT_FILE, data)
-    return items[index]
+    return update_item(kind, text, lambda item: with_optional(item, "rating", rating))
+
+
+def set_status(kind: str, text: str, status: str) -> dict:
+    check_status(status)
+    return update_item(kind, text, lambda item: with_optional(item, "status", None if status == "watched" else status))
+
+
+def set_collection(kind: str, text: str, collection: str | None) -> dict:
+    return update_item(kind, text, lambda item: with_optional(item, "collection", collection))