Support 3.10+ UnionType

Author: bdowning

sql_athame/dataclasses.py

@@ -1,5 +1,6 @@
 import datetime
 import functools
+import sys
 import uuid
 from collections.abc import AsyncGenerator, Iterable, Mapping
 from dataclasses import Field, InitVar, dataclass, fields
@@ -105,12 +106,18 @@ def maybe_serialize(self, value: Any) -> Any:
         return value
 
 
+UNION_TYPES: tuple = (Union,)
+if sys.version_info >= (3, 10):
+    from types import UnionType
+
+    UNION_TYPES = (Union, UnionType)
+
 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:
+    if get_origin(typ) in UNION_TYPES:
         args = []
         for arg in get_args(typ):
             if arg in NULLABLE_TYPES:

tests/test_dataclasses.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import sys
 import uuid
 from dataclasses import dataclass
 from typing import Annotated, Any, Optional, Union
@@ -101,6 +102,27 @@ class Test(ModelBase, table_name="table", primary_key="foo"):
     ]
 
 
+@pytest.mark.skipif(sys.version_info < (3, 10), reason="needs python3.10 or greater")
+def test_py310_unions():
+    @dataclass
+    class Test(ModelBase, table_name="table", primary_key="foo"):
+        foo: int
+        bar: str
+        baz: uuid.UUID | None
+        foo_nullable: int | None
+        bar_nullable: str | None
+
+    assert list(Test.create_table_sql()) == [
+        'CREATE TABLE IF NOT EXISTS "table" ('
+        '"foo" INTEGER NOT NULL, '
+        '"bar" TEXT NOT NULL, '
+        '"baz" UUID, '
+        '"foo_nullable" INTEGER, '
+        '"bar_nullable" TEXT, '
+        'PRIMARY KEY ("foo"))'
+    ]
+
+
 def test_modelclass_missing_type():
     @dataclass
     class Test(ModelBase, table_name="table", primary_key="foo"):