Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Accept and merge multiple ColumnInfo, make type optional #21

Merged
merged 2 commits into from
May 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .bumpversion.cfg
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
[bumpversion]
current_version = 0.4.0-alpha-6
current_version = 0.4.0-alpha-7
parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(-(?P<release>.*)-(?P<build>\d+))?
serialize =
{major}.{minor}.{patch}-{release}-{build}
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "sql-athame"
version = "0.4.0-alpha-6"
version = "0.4.0-alpha-7"
description = "Python tool for slicing and dicing SQL"
authors = ["Brian Downing <bdowning@lavos.net>"]
license = "MIT"
Expand Down
60 changes: 47 additions & 13 deletions sql_athame/dataclasses.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import datetime
import functools
import uuid
from collections.abc import AsyncGenerator, Iterable, Mapping
from dataclasses import Field, InitVar, dataclass, fields
Expand Down Expand Up @@ -29,27 +30,56 @@

@dataclass
class ColumnInfo:
type: str
create_type: str = ""
nullable: bool = False
type: Optional[str] = None
create_type: Optional[str] = None
nullable: Optional[bool] = None
_constraints: tuple[str, ...] = ()

constraints: InitVar[Union[str, Iterable[str], None]] = 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)
if constraints is not None:
if type(constraints) is str:
constraints = (constraints,)
self._constraints = tuple(constraints)

@staticmethod
def merge(a: "ColumnInfo", b: "ColumnInfo") -> "ColumnInfo":
return ColumnInfo(
type=b.type if b.type is not None else a.type,
create_type=b.create_type if b.create_type is not None else a.create_type,
nullable=b.nullable if b.nullable is not None else a.nullable,
_constraints=(*a._constraints, *b._constraints),
)


@dataclass
class ConcreteColumnInfo:
type: str
create_type: str
nullable: bool
constraints: tuple[str, ...]

@staticmethod
def from_column_info(name: str, *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}")
return ConcreteColumnInfo(
type=info.type,
create_type=info.create_type,
nullable=bool(info.nullable),
constraints=info._constraints,
)

def create_table_string(self) -> str:
parts = (
self.create_type,
*(() if self.nullable else ("NOT NULL",)),
*self._constraints,
*self.constraints,
)
return " ".join(parts)

Expand Down Expand Up @@ -86,7 +116,7 @@ def create_table_string(self) -> str:


class ModelBase:
_column_info: Optional[dict[str, ColumnInfo]]
_column_info: Optional[dict[str, ConcreteColumnInfo]]
_cache: dict[tuple, Any]
table_name: str
primary_key_names: tuple[str, ...]
Expand Down Expand Up @@ -138,19 +168,23 @@ def type_hints(cls) -> dict[str, type]:
return cls._type_hints

@classmethod
def column_info_for_field(cls, field: Field) -> ColumnInfo:
def column_info_for_field(cls, field: Field) -> ConcreteColumnInfo:
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
info = []
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):
return md
type, nullable = sql_type_map[base_type]
return ColumnInfo(type=type, nullable=nullable)
info.append(md)
return ConcreteColumnInfo.from_column_info(field.name, *info)

@classmethod
def column_info(cls, column: str) -> ColumnInfo:
def column_info(cls, column: str) -> ConcreteColumnInfo:
try:
return cls._column_info[column] # type: ignore
except AttributeError:
Expand Down
19 changes: 19 additions & 0 deletions tests/test_dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
from dataclasses import dataclass
from typing import Annotated, Optional

import pytest

from sql_athame import sql
from sql_athame.dataclasses import ColumnInfo, ModelBase

Expand Down Expand Up @@ -67,16 +69,33 @@ class Test(ModelBase, table_name="table", primary_key="foo"):
foo: int
bar: str
baz: Optional[uuid.UUID]
quux: Annotated[int, ColumnInfo(constraints="REFERENCES foobar")]
quuux: Annotated[
int,
ColumnInfo(constraints="REFERENCES foobar"),
ColumnInfo(constraints="BLAH", nullable=True),
]

assert list(Test.create_table_sql()) == [
'CREATE TABLE IF NOT EXISTS "table" ('
'"foo" INTEGER NOT NULL, '
'"bar" TEXT NOT NULL, '
'"baz" UUID, '
'"quux" INTEGER NOT NULL REFERENCES foobar, '
'"quuux" INTEGER REFERENCES foobar BLAH, '
'PRIMARY KEY ("foo"))'
]


def test_modelclass_missing_type():
@dataclass
class Test(ModelBase, table_name="table", primary_key="foo"):
foo: dict

with pytest.raises(ValueError, match="Missing SQL type for column 'foo'"):
Test.create_table_sql()


def test_upsert():
@dataclass
class Test(ModelBase, table_name="table", primary_key="id"):
Expand Down