feat: comment dataclass & fix: consistently use CommentSource

Author: faretek1

scratchattach/site/comment.py

@@ -2,103 +2,73 @@
 from __future__ import annotations
 
 import warnings
+import html
+
+from dataclasses import dataclass
+from typing_extensions import assert_never
 from typing import Union, Optional, Any
-from typing_extensions import assert_never  # importing from typing caused me errors
 from enum import Enum, auto
-import html
 
-from . import user, project, studio
+from . import user, project, studio, session
 from ._base import BaseSiteComponent
 from scratchattach.utils import exceptions
 
 class CommentSource(Enum):
     PROJECT = auto()
     USER_PROFILE = auto()
     STUDIO = auto()
+    UNKNOWN = auto()
 
+@dataclass
 class Comment(BaseSiteComponent):
     """
     Represents a Scratch comment (on a profile, studio or project)
     """
-    id: Union[int, str]
-    source: CommentSource
-    source_id: Union[int, str]
-    cached_replies: Optional[list[Comment]]
-    parent_id: Optional[Union[int, str]]
-    cached_parent_comment: Optional[Comment]
-    commentee_id: Optional[int]
-    content: Any
+    id: Optional[int | str] = None
+    source: CommentSource = CommentSource.UNKNOWN
+    source_id: Optional[int | str] = None
+    cached_replies: Optional[list[Comment]] = None
+    parent_id: Optional[int | str] = None
+    cached_parent_comment: Optional[Comment] = None
+    commentee_id: Optional[int] = None
+    content: Optional[str] = None
+    reply_count: Optional[int] = None
+    written_by_scratchteam: Optional[bool] = None
+    author_id: Optional[int] = None
+    author_name: Optional[str] = None
+
+    _session: Optional[session.Session] = None
 
     def __str__(self):
-        return str(self.content)
-
-    def __init__(self, **entries):
-
-        # Set attributes every Comment object needs to have:
-        self.id = None
-        self._session = None
-        self.source = None
-        self.source_id = None
-        self.cached_replies = None
-        self.parent_id = None
-        self.cached_parent_comment = None
-
-        # Update attributes from entries dict:
-        self.__dict__.update(entries)
-
-        if "source" not in entries:
-            self.source = "Unknown"
+        return self.text
 
     def update(self):
         warnings.warn("Warning: Comment objects can't be updated")
         return False  # Objects of this type cannot be updated
 
-    def _update_from_dict(self, data):
-        try:
-            self.id = data["id"]
-        except Exception:
-            pass
-        try:
-            self.parent_id = data["parent_id"]
-        except Exception:
-            pass
-        try:
-            self.commentee_id = data["commentee_id"]
-        except Exception:
-            pass
-        try:
-            self.content = data["content"]
-        except Exception:
-            pass
-        try:
-            self.datetime_created = data["datetime_created"]
-        except Exception:
-            pass
-        try:
-            self.author_name = data["author"]["username"]
-        except Exception:
-            pass
-        try:
-            self.author_id = data["author"]["id"]
-        except Exception:
-            pass
-        try:
-            self.written_by_scratchteam = data["author"]["scratchteam"]
-        except Exception:
-            pass
-        try:
-            self.reply_count = data["reply_count"]
-        except Exception:
-            pass
-        try:
-            self.source = data["source"]
-        except Exception:
-            pass
-        try:
-            self.source_id = data["source_id"]
-        except Exception:
-            pass
-        return True
+    def _update_from_dict(self, data: dict[str, str | dict | Any]):
+        self.id = data["id"]
+        self.parent_id = data.get("parent_id")
+        self.commentee_id = data.get("commentee_id")
+        self.content = str(data["content"])
+        self.datetime_created = data["datetime_created"]
+
+        author = data.get("author", {})
+        self.author_name = author.get("username", self.author_name)
+        self.author_id = author.get("id", self.author_id)
+        self.written_by_scratchteam = author.get("scratchteam", self.written_by_scratchteam)
+        self.reply_count = data.get("reply_count", self.reply_count)
+
+        source: str = data.get("source")
+        if self.source is CommentSource.UNKNOWN:
+            self.source = {
+                "project": CommentSource.PROJECT,
+                "studio": CommentSource.STUDIO,
+                "profile": CommentSource.USER_PROFILE,
+                None: CommentSource.UNKNOWN,
+            }[source]
+
+        self.source_id = data.get("source_id", self.source_id)
 
     @property
     def text(self) -> str:
@@ -109,7 +79,7 @@ def text(self) -> str:
             # user profile comments do not seem to be escaped
             return self.content
 
-        return html.unescape(self.content)
+        return str(html.unescape(self.content))
 
     # Methods for getting related entities
 
@@ -159,16 +129,16 @@ def replies(self, *, use_cache: bool = True, limit=40, offset=0):
             use_cache (bool): Returns the replies cached on the first reply fetch. This makes it SIGNIFICANTLY faster for profile comments. Warning: For profile comments, the replies are retrieved and cached on object creation.
         """
         if (self.cached_replies is None) or (not use_cache):
-            if self.source == "profile":
+            if self.source == CommentSource.USER_PROFILE:
                 self.cached_replies = user.User(username=self.source_id, _session=self._session).comment_by_id(
                     self.id).cached_replies[offset:offset + limit]
 
-            elif self.source == "project":
+            elif self.source == CommentSource.PROJECT:
                 p = project.Project(id=self.source_id, _session=self._session)
                 p.update()
                 self.cached_replies = p.comment_replies(comment_id=self.id, limit=limit, offset=offset)
 
-            elif self.source == "studio":
+            elif self.source == CommentSource.STUDIO:
                 self.cached_replies = studio.Studio(id=self.source_id, _session=self._session).comment_replies(
                     comment_id=self.id, limit=limit, offset=offset)
 
@@ -194,58 +164,63 @@ def reply(self, content, *, commentee_id=None):
 
         Returns:
             scratchattach.Comment: The created comment.
+            :param content: Content of the comment to send
+            :param commentee_id: ID of user to reply to
         """
 
         self._assert_auth()
         parent_id = str(self.id)
         if self.parent_id is not None:
             parent_id = str(self.parent_id)
         if commentee_id is None:
-            if "author_id" in self.__dict__:
+            if self.author_id:
                 commentee_id = self.author_id
             else:
                 commentee_id = ""
-        if self.source == "profile":
+
+        if self.source == CommentSource.USER_PROFILE:
             return user.User(username=self.source_id, _session=self._session).reply_comment(content,
                                                                                             parent_id=str(parent_id),
                                                                                             commentee_id=commentee_id)
-        if self.source == "project":
+        if self.source == CommentSource.PROJECT:
             p = project.Project(id=self.source_id, _session=self._session)
             p.update()
             return p.reply_comment(content, parent_id=str(parent_id), commentee_id=commentee_id)
-        if self.source == "studio":
+
+        if self.source == CommentSource.STUDIO:
             return studio.Studio(id=self.source_id, _session=self._session).reply_comment(content,
                                                                                           parent_id=str(parent_id),
                                                                                           commentee_id=commentee_id)
+        raise ValueError(f"Unknown source: {self.source}")
 
     def delete(self):
         """
         Deletes the comment.
         """
         self._assert_auth()
-        if self.source == "profile":
+        if self.source == CommentSource.USER_PROFILE:
             user.User(username=self.source_id, _session=self._session).delete_comment(comment_id=self.id)
 
-        elif self.source == "project":
+        elif self.source == CommentSource.PROJECT:
             p = project.Project(id=self.source_id, _session=self._session)
             p.update()
             p.delete_comment(comment_id=self.id)
 
-        elif self.source == "studio":
+        elif self.source == CommentSource.STUDIO:
             studio.Studio(id=self.source_id, _session=self._session).delete_comment(comment_id=self.id)
 
     def report(self):
         """
         Reports the comment to the Scratch team.
         """
         self._assert_auth()
-        if self.source == "profile":
+        if self.source == CommentSource.USER_PROFILE:
             user.User(username=self.source_id, _session=self._session).report_comment(comment_id=self.id)
 
-        elif self.source == "project":
+        elif self.source == CommentSource.PROJECT:
             p = project.Project(id=self.source_id, _session=self._session)
             p.update()
             p.report_comment(comment_id=self.id)
 
-        elif self.source == "studio":
+        elif self.source == CommentSource.STUDIO:
             studio.Studio(id=self.source_id, _session=self._session).report_comment(comment_id=self.id)