Merge pull request #22 from bdowning/updates

Try to detect nullability from Python types

Author: bdowning

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.4.0-alpha-7
+current_version = 0.4.0-alpha-8
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(-(?P<release>.*)-(?P<build>\d+))?
 serialize = 
 	{major}.{minor}.{patch}-{release}-{build}

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sql-athame"
-version = "0.4.0-alpha-7"
+version = "0.4.0-alpha-8"
 description = "Python tool for slicing and dicing SQL"
 authors = ["Brian Downing <bdowning@lavos.net>"]
 license = "MIT"

sql_athame/dataclasses.py

@@ -10,6 +10,7 @@
     Optional,
     TypeVar,
     Union,
+    get_args,
     get_origin,
     get_type_hints,
 )
@@ -84,30 +85,38 @@ def create_table_string(self) -> str:
         return " ".join(parts)
 
 
+NULLABLE_TYPES = (type(None), Any, object)
+
+
+def split_nullable(typ: type) -> tuple[bool, type]:
+    nullable = typ in NULLABLE_TYPES
+    if get_origin(typ) is Union:
+        args = []
+        for arg in get_args(typ):
+            if arg in NULLABLE_TYPES:
+                nullable = True
+            else:
+                args.append(arg)
+        return nullable, Union[tuple(args)]  # type: ignore
+    return nullable, typ
+
+
 sql_create_type_map = {
     "BIGSERIAL": "BIGINT",
     "SERIAL": "INTEGER",
     "SMALLSERIAL": "SMALLINT",
 }
 
 
-sql_type_map: dict[Any, tuple[str, bool]] = {
-    Optional[bool]: ("BOOLEAN", True),
-    Optional[bytes]: ("BYTEA", True),
-    Optional[datetime.date]: ("DATE", True),
-    Optional[datetime.datetime]: ("TIMESTAMP", True),
-    Optional[float]: ("DOUBLE PRECISION", True),
-    Optional[int]: ("INTEGER", True),
-    Optional[str]: ("TEXT", True),
-    Optional[uuid.UUID]: ("UUID", True),
-    bool: ("BOOLEAN", False),
-    bytes: ("BYTEA", False),
-    datetime.date: ("DATE", False),
-    datetime.datetime: ("TIMESTAMP", False),
-    float: ("DOUBLE PRECISION", False),
-    int: ("INTEGER", False),
-    str: ("TEXT", False),
-    uuid.UUID: ("UUID", False),
+sql_type_map: dict[Any, str] = {
+    bool: "BOOLEAN",
+    bytes: "BYTEA",
+    datetime.date: "DATE",
+    datetime.datetime: "TIMESTAMP",
+    float: "DOUBLE PRECISION",
+    int: "INTEGER",
+    str: "TEXT",
+    uuid.UUID: "UUID",
 }
 
 
@@ -171,16 +180,16 @@ def type_hints(cls) -> dict[str, type]:
     def column_info_for_field(cls, field: Field) -> ConcreteColumnInfo:
         type_info = cls.type_hints()[field.name]
         base_type = type_info
+        metadata = []
         if get_origin(type_info) is Annotated:
-            base_type = type_info.__origin__  # type: ignore
-        info = []
+            base_type, *metadata = get_args(type_info)
+        nullable, base_type = split_nullable(base_type)
+        info = [ColumnInfo(nullable=nullable)]
         if base_type in sql_type_map:
-            _type, nullable = sql_type_map[base_type]
-            info.append(ColumnInfo(type=_type, nullable=nullable))
-        if get_origin(type_info) is Annotated:
-            for md in type_info.__metadata__:  # type: ignore
-                if isinstance(md, ColumnInfo):
-                    info.append(md)
+            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)
 
     @classmethod

tests/test_dataclasses.py

@@ -4,7 +4,7 @@
 
 import uuid
 from dataclasses import dataclass
-from typing import Annotated, Optional
+from typing import Annotated, Any, Optional, Union
 
 import pytest
 
@@ -75,6 +75,13 @@ class Test(ModelBase, table_name="table", primary_key="foo"):
             ColumnInfo(constraints="REFERENCES foobar"),
             ColumnInfo(constraints="BLAH", nullable=True),
         ]
+        any: Annotated[Any, ColumnInfo(type="TEXT")]
+        any_not_null: Annotated[Any, ColumnInfo(type="TEXT", nullable=False)]
+        obj: Annotated[object, ColumnInfo(type="TEXT")]
+        obj_not_null: Annotated[object, ColumnInfo(type="TEXT", nullable=False)]
+        combined_nullable: Annotated[Union[int, Any], ColumnInfo(type="INTEGER")]
+        null_jsonb: Annotated[Optional[dict], ColumnInfo(type="JSONB")]
+        not_null_jsonb: Annotated[dict, ColumnInfo(type="JSONB")]
 
     assert list(Test.create_table_sql()) == [
         'CREATE TABLE IF NOT EXISTS "table" ('
@@ -83,6 +90,13 @@ class Test(ModelBase, table_name="table", primary_key="foo"):
         '"baz" UUID, '
         '"quux" INTEGER NOT NULL REFERENCES foobar, '
         '"quuux" INTEGER REFERENCES foobar BLAH, '
+        '"any" TEXT, '
+        '"any_not_null" TEXT NOT NULL, '
+        '"obj" TEXT, '
+        '"obj_not_null" TEXT NOT NULL, '
+        '"combined_nullable" INTEGER, '
+        '"null_jsonb" JSONB, '
+        '"not_null_jsonb" JSONB NOT NULL, '
         'PRIMARY KEY ("foo"))'
     ]