Column info via annotation

Author: bdowning

sql_athame/__init__.py

@@ -1,2 +1,2 @@
 from .base import Fragment, sql
-from .dataclasses import ModelBase, model_field, model_field_metadata
+from .dataclasses import ColumnInfo, ModelBase

sql_athame/dataclasses.py

@@ -1,13 +1,16 @@
 import datetime
 import uuid
 from collections.abc import AsyncGenerator, Iterable, Mapping
-from dataclasses import InitVar, dataclass, field, fields
+from dataclasses import Field, InitVar, dataclass, fields
 from typing import (
+    Annotated,
     Any,
     Callable,
     Optional,
     TypeVar,
     Union,
+    get_origin,
+    get_type_hints,
 )
 
 from .base import Fragment, sql
@@ -31,7 +34,7 @@ class ColumnInfo:
 
     constraints: InitVar[Union[str, Iterable[str], None]] = None
 
-    def __post_init__(self, constraints: Union[str, Iterable[str], None]):
+    def __post_init__(self, constraints: Union[str, Iterable[str], None]) -> None:
         if self.create_type == "":
             self.create_type = self.type
             self.type = sql_create_type_map.get(self.type.upper(), self.type)
@@ -40,7 +43,7 @@ def __post_init__(self, constraints: Union[str, Iterable[str], None]):
                 constraints = (constraints,)
             self._constraints = tuple(constraints)
 
-    def create_table_string(self):
+    def create_table_string(self) -> str:
         parts = (
             self.create_type,
             *(() if self.nullable else ("NOT NULL",)),
@@ -49,30 +52,14 @@ def create_table_string(self):
         return " ".join(parts)
 
 
-def model_field_metadata(
-    type: str, nullable: bool = False, constraints: Union[str, Iterable[str]] = ()
-) -> dict[str, Any]:
-    if isinstance(constraints, str):
-        constraints = (constraints,)
-    info = ColumnInfo(type=type, nullable=nullable, constraints=constraints)
-
-    return {"sql_athame": info}
-
-
-def model_field(
-    *, type: str, constraints: Union[str, Iterable[str]] = (), **kwargs: Any
-) -> Any:
-    return field(**kwargs, metadata=model_field_metadata(type, constraints))
-
-
 sql_create_type_map = {
     "BIGSERIAL": "BIGINT",
     "SERIAL": "INTEGER",
     "SMALLSERIAL": "SMALLINT",
 }
 
 
-sql_type_map: dict[type, tuple[str, bool]] = {
+sql_type_map: dict[Any, tuple[str, bool]] = {
     Optional[bool]: ("BOOLEAN", True),
     Optional[bytes]: ("BYTEA", True),
     Optional[datetime.date]: ("DATE", True),
@@ -92,13 +79,6 @@ def model_field(
 }
 
 
-def column_info_for_field(field):
-    if "sql_athame" in field.metadata:
-        return field.metadata["sql_athame"]
-    type, nullable = sql_type_map[field.type]
-    return ColumnInfo(type=type, nullable=nullable)
-
-
 T = TypeVar("T", bound="ModelBase")
 U = TypeVar("U")
 
@@ -109,6 +89,7 @@ class ModelBase:
     table_name: str
     primary_key_names: tuple[str, ...]
     array_safe_insert: bool
+    _type_hints: dict[str, type]
 
     def __init_subclass__(
         cls,
@@ -146,12 +127,34 @@ def _cached(cls, key: tuple, thunk: Callable[[], U]) -> U:
             cls._cache[key] = thunk()
             return cls._cache[key]
 
+    @classmethod
+    def type_hints(cls) -> dict[str, type]:
+        try:
+            return cls._type_hints
+        except AttributeError:
+            cls._type_hints = get_type_hints(cls, include_extras=True)
+            return cls._type_hints
+
+    @classmethod
+    def column_info_for_field(cls, field: Field) -> ColumnInfo:
+        type_info = cls.type_hints()[field.name]
+        base_type = type_info
+        if get_origin(type_info) is Annotated:
+            base_type = type_info.__origin__  # type: ignore
+            for md in type_info.__metadata__:  # type: ignore
+                if isinstance(md, ColumnInfo):
+                    return md
+        type, nullable = sql_type_map[base_type]
+        return ColumnInfo(type=type, nullable=nullable)
+
     @classmethod
     def column_info(cls, column: str) -> ColumnInfo:
         try:
             return cls._column_info[column]  # type: ignore
         except AttributeError:
-            cls._column_info = {f.name: column_info_for_field(f) for f in cls._fields()}
+            cls._column_info = {
+                f.name: cls.column_info_for_field(f) for f in cls._fields()
+            }
             return cls._column_info[column]
 
     @classmethod

tests/test_asyncpg.py

@@ -3,12 +3,13 @@
 import os
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Optional
+from typing import Annotated, Optional
 
 import asyncpg
 import pytest
+from typing_extensions import TypeAlias
 
-from sql_athame import ModelBase, model_field_metadata, sql
+from sql_athame import ColumnInfo, ModelBase, sql
 
 
 @pytest.fixture(autouse=True)
@@ -137,7 +138,7 @@ class Test(
         insert_multiple_mode="array_safe",
     ):
         id: int
-        a: list[int] = field(metadata=model_field_metadata(type="INT[]"))
+        a: Annotated[list[int], ColumnInfo(type="INT[]")]
         b: str
 
     await conn.execute(*Test.create_table_sql())
@@ -260,10 +261,13 @@ class Test(ModelBase, table_name="test", primary_key=("id1", "id2")):
     ]
 
 
+Serial: TypeAlias = Annotated[int, ColumnInfo(type="SERIAL")]
+
+
 async def test_serial(conn):
     @dataclass
     class Test(ModelBase, table_name="table", primary_key="id"):
-        id: int = field(metadata=model_field_metadata(type="SERIAL"))
+        id: Serial
         foo: int
         bar: str
 
@@ -279,10 +283,8 @@ class Test(ModelBase, table_name="table", primary_key="id"):
 async def test_unnest_json(conn):
     @dataclass
     class Test(ModelBase, table_name="table", primary_key="id"):
-        id: int = field(metadata=model_field_metadata(type="SERIAL"))
-        json: Optional[list] = field(
-            metadata=model_field_metadata(type="JSONB", nullable=True)
-        )
+        id: Serial
+        json: Annotated[Optional[list], ColumnInfo(type="JSONB", nullable=True)]
 
     await conn.set_type_codec(
         "jsonb", encoder=json.dumps, decoder=json.loads, schema="pg_catalog"
@@ -307,7 +309,7 @@ class Test(ModelBase, table_name="table", primary_key="id"):
 async def test_unnest_empty(conn):
     @dataclass
     class Test(ModelBase, table_name="table", primary_key="id"):
-        id: int = field(metadata=model_field_metadata(type="SERIAL"))
+        id: Serial
 
     await conn.execute(*Test.create_table_sql())
 

tests/test_dataclasses.py

@@ -1,16 +1,16 @@
 import uuid
-from dataclasses import dataclass, field
-from typing import Optional
+from dataclasses import dataclass
+from typing import Annotated, Optional
 
 from sql_athame import sql
-from sql_athame.dataclasses import ModelBase, model_field_metadata
+from sql_athame.dataclasses import ColumnInfo, ModelBase
 
 
 def test_modelclass():
     @dataclass
     class Test(ModelBase, table_name="table"):
-        foo: int = field(metadata=model_field_metadata(type="INTEGER"))
-        bar: str = field(default="hi", metadata=model_field_metadata(type="TEXT"))
+        foo: int
+        bar: str = "hi"
 
     t = Test(42)
 
@@ -94,7 +94,7 @@ class Test(ModelBase, table_name="table", primary_key="id"):
 def test_serial():
     @dataclass
     class Test(ModelBase, table_name="table", primary_key="id"):
-        id: int = field(metadata=model_field_metadata(type="SERIAL"))
+        id: Annotated[int, ColumnInfo(type="SERIAL")]
         foo: int
         bar: str