Merge pull request #35 from piercefreeman/feature/expand-docstrings

Expand class docstrings with typehints

Author: piercefreeman

README.md

@@ -238,7 +238,3 @@ a simple poetry install. Poetry is set up to create a dynamic `setup.py` based o
 ```bash
 poetry install
 ```
-
-## TODOs
-
-- [ ] Additional documentation with usage examples.

iceaxe/__tests__/test_session.py

@@ -529,6 +529,28 @@ async def test_refresh(db_connection: DBConnection):
     assert user.name == "Jane Doe"
 
 
+@pytest.mark.asyncio
+async def test_get(db_connection: DBConnection):
+    """
+    Test retrieving a single record by primary key using the get method.
+    """
+    # Create a test user
+    user = UserDemo(name="John Doe", email="[email protected]")
+    await db_connection.insert([user])
+    assert user.id is not None
+
+    # Test successful get
+    retrieved_user = await db_connection.get(UserDemo, user.id)
+    assert retrieved_user is not None
+    assert retrieved_user.id == user.id
+    assert retrieved_user.name == "John Doe"
+    assert retrieved_user.email == "[email protected]"
+
+    # Test get with non-existent ID
+    non_existent = await db_connection.get(UserDemo, 9999)
+    assert non_existent is None
+
+
 @pytest.mark.asyncio
 async def test_db_connection_insert_update_enum(db_connection: DBConnection):
     """

iceaxe/base.py

@@ -16,13 +16,40 @@
 
 @dataclass_transform(kw_only_default=True, field_specifiers=(PydanticField,))
 class DBModelMetaclass(_model_construction.ModelMetaclass):
+    """
+    Metaclass for database model classes that provides automatic field tracking and SQL query generation.
+    Extends Pydantic's model metaclass to add database-specific functionality.
+
+    This metaclass provides:
+    - Automatic field to SQL column mapping
+    - Dynamic field access that returns query-compatible field definitions
+    - Registry tracking of all database model classes
+    - Support for generic model instantiation
+
+    ```python {{sticky: True}}
+    class User(TableBase):  # Uses DBModelMetaclass
+        id: int = Field(primary_key=True)
+        name: str
+        email: str | None
+
+    # Fields can be accessed for queries
+    User.id      # Returns DBFieldClassDefinition
+    User.name    # Returns DBFieldClassDefinition
+
+    # Metaclass handles model registration
+    registered_models = DBModelMetaclass.get_registry()
+    ```
+    """
+
     _registry: list[Type["TableBase"]] = []
-    # {class: kwargs}
     _cached_args: dict[Type["TableBase"], dict[str, Any]] = {}
+    is_constructing: bool = False
 
-    def __new__(
-        mcs, name: str, bases: tuple, namespace: dict[str, Any], **kwargs: Any
-    ) -> type:
+    def __new__(mcs, name, bases, namespace, **kwargs):
+        """
+        Create a new database model class with proper field tracking.
+        Handles registration of the model and processes any table-specific arguments.
+        """
         raw_kwargs = {**kwargs}
 
         mcs.is_constructing = True
@@ -59,8 +86,15 @@ def __new__(
         return cls
 
     def __getattr__(self, key: str) -> Any:
-        # Inspired by the approach in our render logic
-        # https://github.com/piercefreeman/mountaineer/blob/fdda3a58c0fafebb43a58b4f3d410dbf44302fd6/mountaineer/render.py#L252
+        """
+        Provides dynamic access to model fields as query-compatible definitions.
+        When accessing an undefined attribute, checks if it's a model field and returns
+        a DBFieldClassDefinition if it is.
+
+        :param key: The attribute name to access
+        :return: Field definition or raises AttributeError
+        :raises AttributeError: If the attribute doesn't exist and isn't a model field
+        """
         if self.is_constructing:
             return super().__getattr__(key)  # type: ignore
 
@@ -78,16 +112,26 @@ def __getattr__(self, key: str) -> Any:
             raise
 
     @classmethod
-    def get_registry(cls):
+    def get_registry(cls) -> list[Type["TableBase"]]:
+        """
+        Get the set of all registered database model classes.
+
+        :return: Set of registered TableBase classes
+        """
         return cls._registry
 
     @classmethod
-    def _extract_kwarg(cls, kwargs: dict[str, Any], key: str, default: Any = None):
+    def _extract_kwarg(
+        cls, kwargs: dict[str, Any], key: str, default: Any = None
+    ) -> Any:
         """
-        Kwarg extraction that supports standard instantiation and pydantic's approach
-        for Generic models where it hydrates a fully new class in memory with the type
-        annotations set to generic values.
+        Extract a keyword argument from either standard kwargs or pydantic generic metadata.
+        Handles both normal instantiation and pydantic's generic model instantiation.
 
+        :param kwargs: Dictionary of keyword arguments
+        :param key: Key to extract
+        :param default: Default value if key not found
+        :return: Extracted value or default
         """
         if key in kwargs:
             return kwargs.pop(key)
@@ -101,69 +145,186 @@ def _extract_kwarg(cls, kwargs: dict[str, Any], key: str, default: Any = None):
 
     @property
     def model_fields(self) -> dict[str, DBFieldInfo]:  # type: ignore
-        # model_fields must be reimplemented in our custom metaclass, otherwise
-        # clients will get the super typehinting signature when they try
-        # to access Model.model_fields. This overrides the ClassVar typehint
-        # that's placed in the TableBase itself.
+        """
+        Get the dictionary of model fields and their definitions.
+        Overrides the ClassVar typehint from TableBase for proper typing.
+
+        :return: Dictionary of field names to field definitions
+        """
         return super().model_fields  # type: ignore
 
 
 class UniqueConstraint(BaseModel):
+    """
+    Represents a UNIQUE constraint in a database table.
+    Ensures that the specified combination of columns contains unique values across all rows.
+
+    ```python {{sticky: True}}
+    class User(TableBase):
+        email: str
+        tenant_id: int
+
+        table_args = [
+            UniqueConstraint(columns=["email", "tenant_id"])
+        ]
+    ```
+    """
+
     columns: list[str]
+    """
+    List of column names that should have unique values
+    """
 
 
 class IndexConstraint(BaseModel):
+    """
+    Represents an INDEX on one or more columns in a database table.
+    Improves query performance for the specified columns.
+
+    ```python {{sticky: True}}
+    class User(TableBase):
+        email: str
+        last_login: datetime
+
+        table_args = [
+            IndexConstraint(columns=["last_login"])
+        ]
+    ```
+    """
+
     columns: list[str]
+    """
+    List of column names to create an index on
+    """
 
 
 INTERNAL_TABLE_FIELDS = ["modified_attrs"]
 
 
 class TableBase(BaseModel, metaclass=DBModelMetaclass):
+    """
+    Base class for all database table models.
+    Provides the foundation for defining database tables using Python classes with
+    type hints and field definitions.
+
+    Features:
+    - Automatic table name generation from class name
+    - Support for custom table names
+    - Tracking of modified fields for efficient updates
+    - Support for unique constraints and indexes
+    - Integration with Pydantic for validation
+
+    ```python {{sticky: True}}
+    class User(TableBase):
+        # Custom table name (optional)
+        table_name = "users"
+
+        # Fields with types and constraints
+        id: int = Field(primary_key=True)
+        email: str = Field(unique=True)
+        name: str
+        is_active: bool = Field(default=True)
+
+        # Table-level constraints
+        table_args = [
+            UniqueConstraint(columns=["email"]),
+            IndexConstraint(columns=["name"])
+        ]
+
+    # Usage in queries
+    query = select(User).where(User.is_active == True)
+    users = await conn.execute(query)
+    ```
+    """
+
     if TYPE_CHECKING:
         model_fields: ClassVar[dict[str, DBFieldInfo]]  # type: ignore
 
     table_name: ClassVar[str] = PydanticUndefined  # type: ignore
+    """
+    Optional custom name for the table
+    """
+
     table_args: ClassVar[list[UniqueConstraint | IndexConstraint]] = PydanticUndefined  # type: ignore
+    """
+    Table constraints and indexes
+    """
 
     # Private methods
     modified_attrs: dict[str, Any] = Field(default_factory=dict, exclude=True)
+    """
+    Dictionary of modified field values since instantiation or the last clear_modified_attributes() call.
+    Used to construct differential update queries.
+    """
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """
+        Track modified attributes when fields are updated.
+        This allows for efficient database updates by only updating changed fields.
 
-    def __setattr__(self, name, value):
+        :param name: Attribute name
+        :param value: New value
+        """
         if name in self.model_fields:
             self.modified_attrs[name] = value
         super().__setattr__(name, value)
 
     def get_modified_attributes(self) -> dict[str, Any]:
+        """
+        Get the dictionary of attributes that have been modified since instantiation
+        or the last clear_modified_attributes() call.
+
+        :return: Dictionary of modified attribute names and their values
+        """
         return self.modified_attrs
 
-    def clear_modified_attributes(self):
+    def clear_modified_attributes(self) -> None:
+        """
+        Clear the tracking of modified attributes.
+        Typically called after successfully saving changes to the database.
+        """
         self.modified_attrs.clear()
 
     @classmethod
-    def get_table_name(cls):
+    def get_table_name(cls) -> str:
+        """
+        Get the table name for this model.
+        Uses the custom table_name if set, otherwise converts the class name to lowercase.
+
+        :return: Table name to use in SQL queries
+        """
         if cls.table_name == PydanticUndefined:
             return cls.__name__.lower()
         return cls.table_name
 
     @classmethod
-    def get_client_fields(cls):
+    def get_client_fields(cls) -> dict[str, DBFieldInfo]:
+        """
+        Get all fields that should be exposed to clients.
+        Excludes internal fields used for model functionality.
+
+        :return: Dictionary of field names to field definitions
+        """
         return {
             field: info
             for field, info in cls.model_fields.items()
             if field not in INTERNAL_TABLE_FIELDS
         }
 
     @classmethod
-    def select_fields(cls):
+    def select_fields(cls) -> QueryLiteral:
         """
-        Returns a query selectable string that can be used to select all fields
-        from this model. This is the format that needs to be passed to our parser
-        to serialize the raw postgres field values as TableBase objects.
+        Generate a SQL-safe string for selecting all fields from this table.
+        The output format is "{table_name}.{field_name} as {table_name}_{field_name}"
+        for each field, which ensures proper field name resolution in complex queries.
 
-        The exact format is formatted as:
-        "{table_name}.{field_name} as {table_name}_{field_name}".
+        :return: SQL-safe field selection string
 
+        ```python {{sticky: True}}
+        # For a User class with fields 'id' and 'name':
+        User.select_fields()
+        # Returns: '"users"."id" as "users_id", "users"."name" as "users_name"'
+        ```
         """
         table_token = QueryIdentifier(cls.get_table_name())
         select_fields: list[str] = []