supython 0.5.0__py3-none-any.whl

supython/gen/_introspect.py

@@ -0,0 +1,137 @@
+"""Shared Postgres schema introspection queries used by type generators.
+
+Both ``types_py.py`` (Python dataclasses) and ``types_ts.py`` (TypeScript
+``Database`` interface) use these functions to fetch metadata from
+``information_schema`` and ``pg_catalog``.
+"""
+
+import asyncpg
+
+
+async def _fetch_enums(
+    conn: asyncpg.Connection, schemas: list[str]
+) -> dict[tuple[str, str], list[str]]:
+    rows = await conn.fetch(
+        """
+        select n.nspname  as schema,
+               t.typname  as name,
+               e.enumlabel as label
+        from pg_type t
+        join pg_enum e       on e.enumtypid = t.oid
+        join pg_namespace n  on n.oid = t.typnamespace
+        where n.nspname = any($1::text[])
+        order by n.nspname, t.typname, e.enumsortorder
+        """,
+        list(schemas),
+    )
+    out: dict[tuple[str, str], list[str]] = {}
+    for r in rows:
+        out.setdefault((r["schema"], r["name"]), []).append(r["label"])
+    return out
+
+
+async def _fetch_tables(
+    conn: asyncpg.Connection, schemas: list[str]
+) -> list[tuple[str, str]]:
+    rows = await conn.fetch(
+        """
+        select table_schema, table_name
+        from information_schema.tables
+        where table_schema = any($1::text[])
+          and table_type in ('BASE TABLE', 'VIEW')
+        order by table_schema, table_name
+        """,
+        list(schemas),
+    )
+    return [(r["table_schema"], r["table_name"]) for r in rows]
+
+
+async def _fetch_columns(
+    conn: asyncpg.Connection, schemas: list[str]
+) -> dict[tuple[str, str], list[asyncpg.Record]]:
+    rows = await conn.fetch(
+        """
+        select c.table_schema,
+               c.table_name,
+               c.column_name,
+               c.ordinal_position,
+               c.is_nullable,
+               c.data_type,
+               c.udt_schema,
+               c.udt_name,
+               c.column_default,
+               c.is_generated,
+               c.is_identity,
+               e.data_type   as element_data_type,
+               e.udt_schema  as element_udt_schema,
+               e.udt_name    as element_udt_name
+        from information_schema.columns c
+        join information_schema.tables t
+          on t.table_catalog = c.table_catalog
+         and t.table_schema  = c.table_schema
+         and t.table_name    = c.table_name
+        left join information_schema.element_types e
+          on  e.object_catalog            = c.table_catalog
+          and e.object_schema             = c.table_schema
+          and e.object_name               = c.table_name
+          and e.object_type               = case t.table_type
+                                                when 'VIEW' then 'VIEW'
+                                                else 'TABLE'
+                                            end
+          and e.collection_type_identifier = c.dtd_identifier
+        where c.table_schema = any($1::text[])
+          and t.table_type in ('BASE TABLE', 'VIEW')
+        order by c.table_schema, c.table_name, c.ordinal_position
+        """,
+        list(schemas),
+    )
+    grouped: dict[tuple[str, str], list[asyncpg.Record]] = {}
+    for r in rows:
+        grouped.setdefault((r["table_schema"], r["table_name"]), []).append(r)
+    return grouped
+
+
+async def _fetch_relationships(
+    conn: asyncpg.Connection, schemas: list[str]
+) -> dict[tuple[str, str], list[dict[str, any]]]:  # type: ignore[no-any-return]
+    rows = await conn.fetch(
+        """
+        select
+            n.nspname as table_schema,
+            c.relname as table_name,
+            con.conname as constraint_name,
+            (
+                select array_agg(a.attname order by ord)
+                from unnest(con.conkey) with ordinality as t(attnum, ord)
+                join pg_attribute a
+                    on a.attnum = t.attnum and a.attrelid = con.conrelid
+            ) as columns,
+            ref_c.relname as foreign_table_name,
+            (
+                select array_agg(a.attname order by ord)
+                from unnest(con.confkey) with ordinality as t(attnum, ord)
+                join pg_attribute a
+                    on a.attnum = t.attnum and a.attrelid = con.confrelid
+            ) as foreign_columns
+        from pg_constraint con
+        join pg_class c on c.oid = con.conrelid
+        join pg_namespace n on n.oid = c.relnamespace
+        join pg_class ref_c on ref_c.oid = con.confrelid
+        where con.contype = 'f'
+          and n.nspname = any($1::text[])
+        order by n.nspname, c.relname, con.conname
+        """,
+        list(schemas),
+    )
+    out: dict[tuple[str, str], list[dict[str, any]]] = {}  # type: ignore[no-any-return]
+    for r in rows:
+        key = (r["table_schema"], r["table_name"])
+        out.setdefault(key, []).append(
+            {
+                "foreignKeyName": r["constraint_name"],
+                "columns": r["columns"],
+                "referencedRelation": r["foreign_table_name"],
+                "referencedColumns": r["foreign_columns"],
+            }
+        )
+    return out

supython/gen/types_py.py

@@ -0,0 +1,270 @@
+"""Render a typed ``types.py`` from Postgres schema introspection.
+
+Public entry point: :func:`render_types_py`. The returned string is the full
+contents of a self-contained Python module with one ``StrEnum`` per Postgres
+enum and one ``@dataclass`` + matching ``TypedDict`` per table or view in the
+selected schemas.
+
+The generator opens its own asyncpg connection and closes it before returning;
+it does not depend on the supython service running.
+"""
+
+import keyword
+import re
+from datetime import UTC, datetime
+
+import asyncpg
+
+from ..settings import get_settings
+from ._introspect import _fetch_columns, _fetch_enums, _fetch_tables
+
+_SIMPLE_MAP: dict[str, tuple[str, str | None]] = {
+    "text": ("str", None),
+    "varchar": ("str", None),
+    "bpchar": ("str", None),
+    "char": ("str", None),
+    "citext": ("str", None),
+    "name": ("str", None),
+    "int2": ("int", None),
+    "int4": ("int", None),
+    "int8": ("int", None),
+    "float4": ("float", None),
+    "float8": ("float", None),
+    "numeric": ("Decimal", "Decimal"),
+    "money": ("Decimal", "Decimal"),
+    "bool": ("bool", None),
+    "uuid": ("UUID", "UUID"),
+    "timestamptz": ("datetime", "datetime"),
+    "timestamp": ("datetime", "datetime"),
+    "date": ("date", "date"),
+    "time": ("time", "time"),
+    "timetz": ("time", "time"),
+    "interval": ("timedelta", "timedelta"),
+    "bytea": ("bytes", None),
+    "json": ("dict[str, Any]", "Any"),
+    "jsonb": ("dict[str, Any]", "Any"),
+    "inet": ("str", None),
+    "cidr": ("str", None),
+    "macaddr": ("str", None),
+    "tsvector": ("str", None),
+    "tsquery": ("str", None),
+    "bit": ("str", None),
+    "varbit": ("str", None),
+    "oid": ("int", None),
+}
+
+
+def _safe_col_name(name: str) -> str:
+    if keyword.iskeyword(name) or keyword.issoftkeyword(name):
+        return f"{name}_"
+    return name
+
+
+async def render_types_py(schemas: list[str]) -> str:
+    """Connect to ``DATABASE_URL`` and return a rendered ``types.py`` module."""
+    s = get_settings()
+    conn = await asyncpg.connect(s.database_url)
+    try:
+        enums = await _fetch_enums(conn, schemas)
+        tables = await _fetch_tables(conn, schemas)
+        columns = await _fetch_columns(conn, schemas)
+    finally:
+        await conn.close()
+    return _render(schemas, enums, tables, columns)
+
+
+
+def _class_name(schema: str, table: str, schemas: list[str]) -> str:
+    base = "".join(part.capitalize() for part in re.split(r"[_\-]", table) if part)
+    if not base:
+        base = "Table"
+    if schema == "public" or len(schemas) == 1:
+        return base
+    return f"{schema.capitalize()}{base}"
+
+
+def _safe_enum_attr(label: str) -> str:
+    s = re.sub(r"[^A-Za-z0-9_]", "_", label).upper()
+    if not s or s[0].isdigit():
+        s = "_" + s
+    return s
+
+
+def _pg_to_py(
+    udt_schema: str,
+    udt_name: str,
+    data_type: str,
+    element: tuple[str, str, str] | None,
+    enum_classes: dict[tuple[str, str], str],
+    imports: set[str],
+) -> tuple[str, str | None]:
+    """Return (annotation, unmapped_comment_or_None)."""
+    if data_type == "ARRAY" and element is not None:
+        elem_ann, elem_unmapped = _pg_to_py(
+            element[1], element[2], element[0], None, enum_classes, imports
+        )
+        return f"list[{elem_ann}]", elem_unmapped
+
+    if data_type == "USER-DEFINED" and (udt_schema, udt_name) in enum_classes:
+        return enum_classes[(udt_schema, udt_name)], None
+
+    if udt_name in _SIMPLE_MAP:
+        ann, imp = _SIMPLE_MAP[udt_name]
+        if imp:
+            imports.add(imp)
+        return ann, None
+
+    imports.add("Any")
+    return "Any", f"unmapped: {udt_schema}.{udt_name}"
+
+
+def _render(
+    schemas: list[str],
+    enums: dict[tuple[str, str], list[str]],
+    tables: list[tuple[str, str]],
+    columns: dict[tuple[str, str], list[asyncpg.Record]],
+) -> str:
+    imports: set[str] = set()
+
+    enum_classes: dict[tuple[str, str], str] = {
+        (schema, name): _class_name(schema, name, schemas)
+        for (schema, name) in enums
+    }
+
+    body: list[str] = []
+
+    for (schema, name), labels in sorted(enums.items()):
+        cls = enum_classes[(schema, name)]
+        body.append("")
+        body.append(f"# --- enum {schema}.{name} {'-' * (60 - len(schema) - len(name))}")
+        body.append("")
+        body.append(f"class {cls}(StrEnum):")
+        for lbl in labels:
+            body.append(f"    {_safe_enum_attr(lbl)} = {lbl!r}")
+
+    has_table = False
+    for schema, table in tables:
+        cols = columns.get((schema, table), [])
+        if not cols:
+            continue
+        has_table = True
+        cls = _class_name(schema, table, schemas)
+
+        rendered_cols: list[tuple[str, str, bool, str | None]] = []
+        for c in cols:
+            element = None
+            if c["element_data_type"]:
+                element = (
+                    c["element_data_type"],
+                    c["element_udt_schema"],
+                    c["element_udt_name"],
+                )
+            ann, unmapped = _pg_to_py(
+                c["udt_schema"],
+                c["udt_name"],
+                c["data_type"],
+                element,
+                enum_classes,
+                imports,
+            )
+            nullable = c["is_nullable"] == "YES"
+            col_name = _safe_col_name(c["column_name"])
+            rendered_cols.append((col_name, ann, nullable, unmapped))
+
+        body.append("")
+        body.append(f"# --- {schema}.{table} {'-' * (64 - len(schema) - len(table))}")
+        body.append("")
+        body.append("@dataclass(kw_only=True, slots=True)")
+        body.append(f"class {cls}:")
+        for col_name, ann, nullable, unmapped in rendered_cols:
+            line = (
+                f"    {col_name}: {ann} | None = None"
+                if nullable
+                else f"    {col_name}: {ann}"
+            )
+            if unmapped:
+                line += f"  # {unmapped}"
+            body.append(line)
+
+        body.append("")
+        body.append("    @classmethod")
+        body.append(f'    def from_record(cls, record: "object") -> "{cls}":')
+        body.append("        fields = cls.__dataclass_fields__")
+        body.append(
+            "        return cls(**{f: v for f, v in record.items() if f in fields})"
+        )
+
+        has_kw = any(
+            _safe_col_name(c["column_name"]) != c["column_name"]
+            for c in cols
+        )
+        if has_kw:
+            _emit_typeddict_functional(body, cls, rendered_cols)
+        else:
+            _emit_typeddict_class(body, cls, rendered_cols)
+
+    header: list[str] = []
+    header.append('"""Generated by `supython gen types --lang py`. Do not edit.')
+    header.append("")
+    header.append(f"Schemas: {', '.join(schemas)}")
+    header.append(f"Generated at: {datetime.now(UTC).isoformat()}")
+    header.append('"""')
+
+    import_lines: list[str] = []
+    if has_table:
+        import_lines.append("from dataclasses import dataclass")
+
+    datetime_syms = sorted(
+        {i for i in imports if i in {"datetime", "date", "time", "timedelta"}}
+    )
+    if datetime_syms:
+        import_lines.append(f"from datetime import {', '.join(datetime_syms)}")
+    if "Decimal" in imports:
+        import_lines.append("from decimal import Decimal")
+    if enums:
+        import_lines.append("from enum import StrEnum")
+    typing_syms: set[str] = set()
+    if "Any" in imports:
+        typing_syms.add("Any")
+    if has_table:
+        typing_syms.add("TypedDict")
+    if typing_syms:
+        import_lines.append(f"from typing import {', '.join(sorted(typing_syms))}")
+    if "UUID" in imports:
+        import_lines.append("from uuid import UUID")
+
+    parts: list[str] = []
+    parts.extend(header)
+    if import_lines:
+        parts.append("")
+        parts.extend(import_lines)
+    parts.extend(body)
+    return "\n".join(parts).rstrip() + "\n"
+
+
+def _emit_typeddict_class(
+    body: list[str],
+    cls: str,
+    rendered_cols: list[tuple[str, str, bool, str | None]],
+) -> None:
+    body.append("")
+    body.append(f"class {cls}Row(TypedDict):")
+    for col_name, ann, nullable, _ in rendered_cols:
+        body.append(
+            f"    {col_name}: {ann} | None" if nullable else f"    {col_name}: {ann}"
+        )
+
+
+def _emit_typeddict_functional(
+    body: list[str],
+    cls: str,
+    rendered_cols: list[tuple[str, str, bool, str | None]],
+) -> None:
+    body.append(f"{cls}Row = TypedDict(\"{cls}Row\", {{")
+    for col_name, ann, nullable, _ in rendered_cols:
+        ann_str = f"{ann} | None" if nullable else ann
+        body.append(f'    "{col_name}": {ann_str},')
+    body.append("})")
+
+
+__all__ = ["render_types_py"]