diff --git a/sqlalchemy_file/base.py b/sqlalchemy_file/base.py index ab34168..a44172e 100644 --- a/sqlalchemy_file/base.py +++ b/sqlalchemy_file/base.py @@ -1,5 +1,20 @@ import typing -from typing import Any +from typing import Any, Literal, overload, override + +STR_ATTRS = Literal[ + "url", + "filename", + "content_type", + "file_id", + "upload_storage", + "uploaded_at", + "path", + "url", +] +BOOL_ATTRS = Literal["saved"] +INT_ATTRS = Literal["size"] +DICT_ATTRS = Literal["meta_data"] +STR_LIST_ATTRS = Literal["files"] class BaseFile(typing.Dict[str, Any]): @@ -11,15 +26,57 @@ class BaseFile(typing.Dict[str, Any]): """ + @overload + def __getitem__(self, key: STR_ATTRS) -> str: + ... + + @overload + def __getitem__(self, key: INT_ATTRS) -> int: + ... + + @overload + def __getitem__(self, key: DICT_ATTRS) -> dict[str, str]: + ... + + @overload + def __getitem__(self, key: STR_LIST_ATTRS) -> list[str]: + ... + + @overload + def __getitem__(self, key: BOOL_ATTRS) -> bool: + ... + + @override def __getitem__(self, key: str) -> Any: return dict.__getitem__(self, key) + @overload + def __getattr__(self, name: STR_ATTRS) -> str: + ... + + @overload + def __getattr__(self, name: INT_ATTRS) -> int: + ... + + @overload + def __getattr__(self, name: DICT_ATTRS) -> dict[str, str]: + ... + + @overload + def __getattr__(self, name: STR_LIST_ATTRS) -> list[str]: + ... + + @overload + def __getattr__(self, name: BOOL_ATTRS) -> bool: + ... + def __getattr__(self, name: str) -> Any: try: return self[name] except KeyError: raise AttributeError(name) + @override def __setitem__(self, key: str, value: Any) -> None: if getattr(self, "_frozen", False): raise TypeError("Already saved files are immutable") @@ -27,6 +84,7 @@ def __setitem__(self, key: str, value: Any) -> None: __setattr__ = __setitem__ + @override def __delattr__(self, name: str) -> None: if getattr(self, "_frozen", False): raise TypeError("Already saved files are immutable") @@ -36,6 +94,7 @@ def __delattr__(self, name: str) -> None: except KeyError: raise AttributeError(name) + @override def __delitem__(self, key: str) -> None: if object.__getattribute__(self, "_frozen"): raise TypeError("Already saved files are immutable")