refactor: add more ruff checks

Author: frostbyte-ninja

cortex_shell/cache.py

@@ -1,5 +1,6 @@
 from collections.abc import Generator
-from hashlib import md5
+from functools import wraps
+from hashlib import sha256
 from pathlib import Path
 from typing import Any, Callable
 
@@ -23,14 +24,15 @@ def __init__(self, size: int, cache_path: Path) -> None:
         self._cache_path = cache_path
         self._cache_path.mkdir(parents=True, exist_ok=True)
 
-    def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:
+    def __call__(self, func: Callable[..., Generator[str, None, None]]) -> Callable[..., Generator[str, None, None]]:
         """
         The Cache decorator.
 
         :param func: The function to cache.
         :return: Wrapped function with caching.
         """
 
+        @wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Generator[str, None, None]:
             cache_file = self._cache_path / self._get_hash_from_request(**kwargs)
             if kwargs.pop("caching", False) and cache_file.exists():
@@ -66,7 +68,7 @@ def _get_hash_from_request(**kwargs: Any) -> str:
         kwargs.pop("caching")
         # delete every message except the last one, which is the most recent user prompt
         kwargs["messages"] = kwargs["messages"][-1:]
-        return md5(yaml_dump_str(kwargs).encode("utf-8")).hexdigest()
+        return sha256(yaml_dump_str(kwargs).encode("utf-8")).hexdigest()
 
     @classmethod
     @option_callback

cortex_shell/error_handler.py

@@ -91,9 +91,9 @@ def _write_line(s: str | None = None, **kwargs: Any) -> None:
         ErrorHandler._write_line_b(s.encode() if s is not None else s, **kwargs)
 
     @staticmethod
-    def _force_bytes(exc: Any) -> bytes:
+    def _force_bytes(obj: Any) -> bytes:  # noqa: ANN401
         with contextlib.suppress(TypeError):
-            return bytes(exc)
+            return bytes(obj)
         with contextlib.suppress(Exception):
-            return str(exc).encode()
-        return f"<unprintable {type(exc).__name__} object>".encode()
+            return str(obj).encode()
+        return f"<unprintable {type(obj).__name__} object>".encode()

cortex_shell/handlers/repl_handler.py

@@ -41,7 +41,7 @@ def _get_prompt_from_input(self, additional_prompt: str | None = None, **kwargs:
 
         super().handle(prompt, **kwargs)
 
-    def _get_user_input(self) -> Any:
+    def _get_user_input(self) -> str:
         bindings = KeyBindings()
 
         @bindings.add(Keys.ControlE)  # type: ignore[misc]

cortex_shell/post_processing/shell_execution_post_processing.py

@@ -31,7 +31,12 @@ class Option(Enum):
 
 
 class ShellExecutionPostProcessing(IPostProcessing):
-    def __init__(self, shell_role: BuiltinRoleShell, describe_shell_role: BuiltinRoleDescribeShell, client: IClient):
+    def __init__(
+        self,
+        shell_role: BuiltinRoleShell,
+        describe_shell_role: BuiltinRoleDescribeShell,
+        client: IClient,
+    ) -> None:
         self._shell_role = shell_role
         self._describe_shell_role = describe_shell_role
         self._handler = self._get_shell_describe_handler(client)

cortex_shell/renderer/plain_renderer.py

@@ -12,7 +12,7 @@
 
 
 class PlainRenderer(IRenderer):
-    def __init__(self, role: Role):
+    def __init__(self, role: Role) -> None:
         self._role = role
 
     def __enter__(self) -> Self:
@@ -22,5 +22,5 @@ def __exit__(self, *args: object) -> None:
         # new line
         print_formatted_text()
 
-    def __call__(self, text: str, chunk: str) -> None:
+    def __call__(self, text: str, chunk: str) -> None:  # noqa: ARG002
         print_formatted_text(get_colored_text(chunk, self._role.output.color), end="")  # type: ignore[arg-type]

cortex_shell/session/chat_session.py

@@ -1,18 +1,20 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
+from ..types import YamlType
 from ..yaml import yaml_dump, yaml_load
 
 if TYPE_CHECKING:  # pragma: no cover
     from pathlib import Path
 
-    from ..configuration.schema import Role
+    from ruamel.yaml import StreamType
+
     from ..types import Message
 
 
 class ChatSession:
-    def __init__(self, file_path: Path, size: int | None = None):
+    def __init__(self, file_path: Path, size: int | None = None) -> None:
         self._file_path = file_path
         self._size = size
 
@@ -38,32 +40,24 @@ def messages(self) -> list[Message]:
     def write_messages(self, messages: list[Message]) -> None:
         if self._size is not None:
             messages = messages[-self._size :]
-        self._write("messages", messages)
-
-    def get_role_name(self) -> str | None:
-        return self._read("role")
-
-    def set_role(self, role: Role) -> None:
-        self._write("role", role.name)
-        self._write("system_prompt", role.description)
-
-    def get_system_prompt(self) -> str:
-        return self._read("system_prompt") or ""
+        self._write("messages", cast(YamlType, messages))
 
-    def _write(self, key: str, value: Any) -> None:
+    def _write(self, key: str, value: YamlType) -> None:
         data = self._yaml_load() if self._file_path.exists() else {}
         data[key] = value
 
         self._yaml_dump(data)
 
-    def _read(self, key: str) -> Any | None:
+    def _read(self, key: str) -> YamlType:
         data = self._yaml_load()
-        return data.get(key) if data else None
+        if isinstance(data, dict):
+            return data.get(key)
+        return None
 
-    def _yaml_load(self) -> Any:
+    def _yaml_load(self) -> Any:  # noqa: ANN401
         if not self._file_path.exists():
             return None
         return yaml_load(stream=self._file_path)
 
-    def _yaml_dump(self, data: Any) -> Any:
+    def _yaml_dump(self, data: Path | StreamType) -> Any:  # noqa: ANN401
         yaml_dump(data=data, stream=self._file_path)

cortex_shell/session/chat_session_manager.py

@@ -14,7 +14,7 @@
 
 
 class ChatSessionManager:
-    def __init__(self, storage_path: Path, history_size: int | None = None):
+    def __init__(self, storage_path: Path, history_size: int | None = None) -> None:
         self._storage_path = storage_path
         self._history_size = history_size
         self._storage_path.mkdir(parents=True, exist_ok=True)

cortex_shell/types/__init__.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
-from typing import Protocol, TypedDict
-
-
-class StringConvertible(Protocol):
-    def __str__(self) -> str:
-        pass
+from typing import TypedDict, Union
 
 
 class Message(TypedDict):
     role: str
     content: str
+
+
+YamlType = Union[str, int, float, bool, None, list["YamlType"], dict[str, "YamlType"]]

cortex_shell/types/prompt_toolkit.py

@@ -39,8 +39,6 @@ def __init__(  # noqa: C901
         values: Sequence[tuple[_T, AnyFormattedText]],
         default_value: _T | None = None,
     ) -> None:
-        assert len(values) > 0
-
         self.values = values
 
         keys: list[_T] = [value for (value, _) in values]

cortex_shell/util.py

@@ -13,11 +13,12 @@
 from importlib import resources
 from pathlib import Path
 from tempfile import gettempdir
-from typing import Any, Callable, cast
+from typing import Any, Callable, Union, cast
 
 import click.exceptions
 import distro
 import typer
+from click import Context
 from pathvalidate import is_valid_filepath
 from prompt_toolkit import print_formatted_text as print_formatted_text_orig
 from prompt_toolkit.formatted_text import FormattedText
@@ -26,12 +27,16 @@
 
 from . import constants as C  # noqa: N812
 
+_TypeOrContext = Union[type, Context]
 
-def option_callback(func: Callable[[Any, str], Any]) -> Callable[[Any, str], Any]:
-    def wrapper(cls: Any, value: str) -> None:
+
+def option_callback(
+    func: Callable[[_TypeOrContext, Any], None],
+) -> Callable[[_TypeOrContext, Any], None]:
+    def wrapper(cls_or_ctx: _TypeOrContext, value: Any) -> None:  # noqa: ANN401
         if not value:
             return
-        func(cls, value)
+        func(cls_or_ctx, value)
         raise typer.Exit
 
     return wrapper
@@ -189,22 +194,13 @@ def get_cache_dir() -> Path:
 
 
 def print_formatted_text(*values: str | FormattedText, **kwargs: Any) -> None:
-    # workaround for NoConsoleScreenBufferError
     if is_tty():
         print_formatted_text_orig(*values, **kwargs)
     else:
-
-        def to_text(val: Any) -> str:
-            if isinstance(val, str):
-                return val
-            elif isinstance(val, list):
-                return "".join([v[1] for v in val])
-            else:
-                raise TypeError
-
-        end = kwargs.get("end", "\n")
-        sep = kwargs.get("sep", " ")
-        print(sep.join([to_text(value) for value in values]), end=end)
+        print(
+            *(text if isinstance(text, str) else "".join(segment[1] for segment in text) for text in values),
+            **kwargs,
+        )
 
 
 def rmtree(path: Path) -> None:

cortex_shell/yaml.py

@@ -10,7 +10,7 @@
 from pydantic import BaseModel
 
 if TYPE_CHECKING:  # pragma: no cover
-    from ruamel.yaml import ScalarNode, StreamTextType, StreamType
+    from ruamel.yaml import BaseRepresenter, ScalarNode, StreamTextType, StreamType
 
 
 class YAML(ruamel.yaml.YAML):
@@ -25,41 +25,41 @@ def __init__(self) -> None:
         self.representer.add_representer(type(None), self._represent_none)
 
     @staticmethod
-    def _represent_none(representer: Any, data: Any) -> ScalarNode:
+    def _represent_none(representer: BaseRepresenter, _data: Any) -> ScalarNode:  # noqa: ANN401
         return representer.represent_scalar("tag:yaml.org,2002:null", "")
 
 
-def yaml_load(stream: Path | StreamTextType) -> Any:
+def yaml_load(stream: Path | StreamTextType) -> Any:  # noqa: ANN401
     if isinstance(stream, Path):
         with stream.open("r", encoding="utf-8") as file:
             return YAML().load(file)
     else:
         return YAML().load(stream)
 
 
-def yaml_dump(data: Path | StreamType, stream: Any = None, *, transform: Any = None) -> Any:
+def yaml_dump(data: Path | StreamType, stream: Any = None, *, transform: Any = None) -> Any:  # noqa: ANN401
     if isinstance(stream, Path):
         with stream.open("w", encoding="utf-8") as file:
-            YAML().dump(data=data, stream=file, transform=transform)
+            return YAML().dump(data=data, stream=file, transform=transform)
     else:
-        YAML().dump(data=data, stream=stream, transform=transform)
+        return YAML().dump(data=data, stream=stream, transform=transform)
 
 
-def yaml_dump_str(data: Path | StreamType, *, transform: Any = None) -> Any:
+def yaml_dump_str(data: Path | StreamType, *, transform: Any = None) -> str:  # noqa: ANN401
     stream = StringIO()
     yaml_dump(data=data, stream=stream, transform=transform)
     return stream.getvalue()
 
 
-T = TypeVar("T", bound=BaseModel)
+_T = TypeVar("_T", bound=BaseModel)
 
 
-def from_yaml_file(model_type: type[T], file: Path) -> T:
+def from_yaml_file(model_type: type[_T], file: Path) -> _T:
     return pydantic.TypeAdapter(model_type).validate_python(yaml_load(file.resolve()))
 
 
 def to_yaml_file(
     file: Path,
-    model: BaseModel | Any,
+    model: BaseModel,
 ) -> None:
     yaml_dump(json.loads(model.model_dump_json()), file.resolve())