Refactor column_info

Author: bdowning

sql_athame/dataclasses.py

@@ -34,8 +34,8 @@ class ColumnInfo:
     type: Optional[str] = None
     create_type: Optional[str] = None
     nullable: Optional[bool] = None
-    _constraints: tuple[str, ...] = ()
 
+    _constraints: tuple[str, ...] = ()
     constraints: InitVar[Union[str, Iterable[str], None]] = None
 
     def __post_init__(self, constraints: Union[str, Iterable[str], None]) -> None:
@@ -56,20 +56,26 @@ def merge(a: "ColumnInfo", b: "ColumnInfo") -> "ColumnInfo":
 
 @dataclass
 class ConcreteColumnInfo:
+    field: Field
+    type_hint: type
     type: str
     create_type: str
     nullable: bool
     constraints: tuple[str, ...]
 
     @staticmethod
-    def from_column_info(name: str, *args: ColumnInfo) -> "ConcreteColumnInfo":
+    def from_column_info(
+        field: Field, type_hint: Any, *args: ColumnInfo
+    ) -> "ConcreteColumnInfo":
         info = functools.reduce(ColumnInfo.merge, args, ColumnInfo())
         if info.create_type is None and info.type is not None:
             info.create_type = info.type
             info.type = sql_create_type_map.get(info.type.upper(), info.type)
         if type(info.type) is not str or type(info.create_type) is not str:
-            raise ValueError(f"Missing SQL type for column {name!r}")
+            raise ValueError(f"Missing SQL type for column {field.name!r}")
         return ConcreteColumnInfo(
+            field=field,
+            type_hint=type_hint,
             type=info.type,
             create_type=info.create_type,
             nullable=bool(info.nullable),
@@ -108,7 +114,7 @@ def split_nullable(typ: type) -> tuple[bool, type]:
 }
 
 
-sql_type_map: dict[Any, str] = {
+sql_type_map: dict[type, str] = {
     bool: "BOOLEAN",
     bytes: "BYTEA",
     datetime.date: "DATE",
@@ -125,12 +131,11 @@ def split_nullable(typ: type) -> tuple[bool, type]:
 
 
 class ModelBase:
-    _column_info: Optional[dict[str, ConcreteColumnInfo]]
+    _column_info: dict[str, ConcreteColumnInfo]
     _cache: dict[tuple, Any]
     table_name: str
     primary_key_names: tuple[str, ...]
     array_safe_insert: bool
-    _type_hints: dict[str, type]
 
     def __init_subclass__(
         cls,
@@ -153,13 +158,6 @@ def __init_subclass__(
         else:
             cls.primary_key_names = tuple(primary_key)
 
-    @classmethod
-    def _fields(cls):
-        # wrapper to ignore typing weirdness: 'Argument 1 to "fields"
-        # has incompatible type "..."; expected "DataclassInstance |
-        # type[DataclassInstance]"'
-        return fields(cls)  # type: ignore
-
     @classmethod
     def _cached(cls, key: tuple, thunk: Callable[[], U]) -> U:
         try:
@@ -169,38 +167,31 @@ def _cached(cls, key: tuple, thunk: Callable[[], U]) -> U:
             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) -> ConcreteColumnInfo:
-        type_info = cls.type_hints()[field.name]
-        base_type = type_info
+    def column_info_for_field(cls, field: Field, type_hint: type) -> ConcreteColumnInfo:
+        base_type = type_hint
         metadata = []
-        if get_origin(type_info) is Annotated:
-            base_type, *metadata = get_args(type_info)
+        if get_origin(type_hint) is Annotated:
+            base_type, *metadata = get_args(type_hint)
         nullable, base_type = split_nullable(base_type)
         info = [ColumnInfo(nullable=nullable)]
         if base_type in sql_type_map:
             info.append(ColumnInfo(type=sql_type_map[base_type]))
         for md in metadata:
             if isinstance(md, ColumnInfo):
                 info.append(md)
-        return ConcreteColumnInfo.from_column_info(field.name, *info)
+        return ConcreteColumnInfo.from_column_info(field, type_hint, *info)
 
     @classmethod
-    def column_info(cls, column: str) -> ConcreteColumnInfo:
+    def column_info(cls) -> dict[str, ConcreteColumnInfo]:
         try:
-            return cls._column_info[column]  # type: ignore
+            return cls._column_info
         except AttributeError:
+            type_hints = get_type_hints(cls, include_extras=True)
             cls._column_info = {
-                f.name: cls.column_info_for_field(f) for f in cls._fields()
+                f.name: cls.column_info_for_field(f, type_hints[f.name])
+                for f in fields(cls)  # type: ignore
             }
-            return cls._column_info[column]
+            return cls._column_info
 
     @classmethod
     def table_name_sql(cls, *, prefix: Optional[str] = None) -> Fragment:
@@ -212,7 +203,11 @@ def primary_key_names_sql(cls, *, prefix: Optional[str] = None) -> list[Fragment
 
     @classmethod
     def field_names(cls, *, exclude: FieldNamesSet = ()) -> list[str]:
-        return [f.name for f in cls._fields() if f.name not in exclude]
+        return [
+            ci.field.name
+            for ci in cls.column_info().values()
+            if ci.field.name not in exclude
+        ]
 
     @classmethod
     def field_names_sql(
@@ -231,9 +226,9 @@ def _get_field_values_fn(
     ) -> Callable[[T], list[Any]]:
         env: dict[str, Any] = {}
         func = ["def get_field_values(self): return ["]
-        for f in cls._fields():
-            if f.name not in exclude:
-                func.append(f"self.{f.name},")
+        for ci in cls.column_info().values():
+            if ci.field.name not in exclude:
+                func.append(f"self.{ci.field.name},")
         func += ["]"]
         exec(" ".join(func), env)
         return env["get_field_values"]
@@ -256,19 +251,15 @@ def field_values_sql(
         else:
             return [sql.value(value) for value in self.field_values()]
 
-    @classmethod
-    def from_tuple(
-        cls: type[T], tup: tuple, *, offset: int = 0, exclude: FieldNamesSet = ()
-    ) -> T:
-        names = (f.name for f in cls._fields() if f.name not in exclude)
-        kwargs = {name: tup[offset] for offset, name in enumerate(names, start=offset)}
-        return cls(**kwargs)
-
     @classmethod
     def from_dict(
         cls: type[T], dct: dict[str, Any], *, exclude: FieldNamesSet = ()
     ) -> T:
-        names = {f.name for f in cls._fields() if f.name not in exclude}
+        names = {
+            ci.field.name
+            for ci in cls.column_info().values()
+            if ci.field.name not in exclude
+        }
         kwargs = {k: v for k, v in dct.items() if k in names}
         return cls(**kwargs)
 
@@ -283,10 +274,10 @@ def create_table_sql(cls) -> Fragment:
         entries = [
             sql(
                 "{} {}",
-                sql.identifier(f.name),
-                sql.literal(cls.column_info(f.name).create_table_string()),
+                sql.identifier(ci.field.name),
+                sql.literal(ci.create_table_string()),
             )
-            for f in cls._fields()
+            for ci in cls.column_info().values()
         ]
         if cls.primary_key_names:
             entries += [sql("PRIMARY KEY ({})", sql.list(cls.primary_key_names_sql()))]
@@ -428,10 +419,11 @@ def delete_multiple_sql(cls: type[T], rows: Iterable[T]) -> Fragment:
                 pks=sql.list(sql.identifier(pk) for pk in cls.primary_key_names),
             ).compile(),
         )
+        column_info = cls.column_info()
         return cached(
             unnest=sql.unnest(
                 (row.primary_key() for row in rows),
-                (cls.column_info(pk).type for pk in cls.primary_key_names),
+                (column_info[pk].type for pk in cls.primary_key_names),
             ),
         )
 
@@ -451,10 +443,11 @@ def insert_multiple_sql(cls: type[T], rows: Iterable[T]) -> Fragment:
                 fields=sql.list(cls.field_names_sql()),
             ).compile(),
         )
+        column_info = cls.column_info()
         return cached(
             unnest=sql.unnest(
                 (row.field_values() for row in rows),
-                (cls.column_info(name).type for name in cls.field_names()),
+                (column_info[name].type for name in cls.field_names()),
             ),
         )
 
@@ -545,9 +538,9 @@ def _get_equal_ignoring_fn(
     ) -> Callable[[T, T], bool]:
         env: dict[str, Any] = {}
         func = ["def equal_ignoring(a, b):"]
-        for f in cls._fields():
-            if f.name not in ignore:
-                func.append(f" if a.{f.name} != b.{f.name}: return False")
+        for ci in cls.column_info().values():
+            if ci.field.name not in ignore:
+                func.append(f" if a.{ci.field.name} != b.{ci.field.name}: return False")
         func += [" return True"]
         exec("\n".join(func), env)
         return env["equal_ignoring"]
@@ -603,9 +596,11 @@ def _get_differences_ignoring_fn(
             "def differences_ignoring(a, b):",
             " diffs = []",
         ]
-        for f in cls._fields():
-            if f.name not in ignore:
-                func.append(f" if a.{f.name} != b.{f.name}: diffs.append({f.name!r})")
+        for ci in cls.column_info().values():
+            if ci.field.name not in ignore:
+                func.append(
+                    f" if a.{ci.field.name} != b.{ci.field.name}: diffs.append({ci.field.name!r})"
+                )
         func += [" return diffs"]
         exec("\n".join(func), env)
         return env["differences_ignoring"]

tests/test_dataclasses.py

@@ -135,8 +135,8 @@ class Test(ModelBase, table_name="table", primary_key="id"):
         foo: int
         bar: str
 
-    assert Test.column_info("id").type == "INTEGER"
-    assert Test.column_info("id").create_type == "SERIAL"
+    assert Test.column_info()["id"].type == "INTEGER"
+    assert Test.column_info()["id"].create_type == "SERIAL"
     assert list(Test.create_table_sql()) == [
         'CREATE TABLE IF NOT EXISTS "table" ('
         '"id" SERIAL NOT NULL, '