sqlty 0.1.0__cp315-cp315-win_amd64.whl

sqlty/_json_shape.py

@@ -0,0 +1,299 @@
+"""Check the JSON a query builds against the annotation of the row field that receives it.
+
+The engine reports the shape of a ``jsonb_build_object()``, ``jsonb_agg()`` or ``to_jsonb()``
+column as JSON text (see ``rust/crates/sqlty-sql/src/json.rs``); drivers decode the value into
+dicts, lists and scalars. `json_problems` compares that shape with the annotation, following
+nested TypedDicts, ``list[...]`` and ``dict[str, ...]``, and lists what cannot hold it.
+
+Only what is decidable is reported: ``Any``, ``object``, classes sqlty cannot see into
+(imports, Pydantic models) and annotations it cannot read accept anything.
+"""
+
+from __future__ import annotations
+
+import ast
+import json
+import typing
+
+if typing.TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from sqlty._type_extractor import TypeInfo
+
+# Annotation heads a JSON scalar decodes into (bool is an int, an int a float; a number
+# is an int when integral, else a float, and ``Decimal`` loads it exactly).
+_SCALAR_HEADS = {
+    "boolean": {"bool", "int", "float"},
+    "integer": {"int", "float"},
+    "number": {"float", "int", "Decimal"},
+    "string": {"str"},
+}
+_SCALAR_NAMES = {
+    "boolean": "a boolean",
+    "integer": "an integer",
+    "number": "a number",
+    "string": "a string",
+}
+_LISTS = {"list", "List", "Sequence", "MutableSequence", "Collection", "Iterable"}
+_MAPPINGS = {"dict", "Dict", "Mapping", "MutableMapping"}
+# Classes that can never hold a JSON value of another kind.
+_BUILTINS = {"bool", "int", "float", "str", "bytes", "tuple", "Tuple", "set", "frozenset"}
+_OPEN = {"Any", "object"}
+_MAX_PROBLEMS = 3
+_MAX_DEPTH = 32
+
+
+class _Context(typing.NamedTuple):
+    defs: Mapping[str, TypeInfo]
+    aliases: Mapping[str, str]
+
+
+def json_problems(
+    shape: str,
+    annotation: ast.expr,
+    defs: Mapping[str, TypeInfo],
+    aliases: Mapping[str, str],
+    name: str = "",
+) -> list[str]:
+    """What keeps ``annotation`` from holding the JSON of ``shape``, as sentences; empty if it fits.
+
+    ``shape`` is the JSON text the engine reports for a column; ``defs`` are the user types
+    of the file (see ``extract_type_definitions``) and ``aliases`` its typing names. Each
+    sentence starts with the path of the value below ``name``: ``items[].price is int but ...``.
+
+    Example:
+        >>> shape = '{"array":{"shape":"integer","nullable":true}}'
+        >>> annotation = ast.parse("list[int]", mode="eval").body
+        >>> json_problems(shape, annotation, {}, {})
+        ['[] can be null but is declared int']
+        >>> json_problems(shape, ast.parse("dict[str, int]", mode="eval").body, {}, {})
+        ['is dict[str, int] but the query builds an array']
+    """
+    out: list[str] = []
+    value = {"shape": json.loads(shape), "nullable": False}
+    _check(value, annotation, name, _Context(defs, aliases), out, 0)
+    return out
+
+
+def render(shape: str) -> str:
+    """The JSON shape the engine reports, as readable text.
+
+    Objects are ``{key: shape}`` (``key?`` when it may be missing), arrays ``[shape]``,
+    maps ``{[key]: shape}``, and a value that can be ``null`` ends in ``| null``.
+
+    Example:
+        >>> render('{"array":{"shape":{"object":[{"name":"id","value":{"shape":"integer","nullable":false},"optional":true}]},"nullable":true}}')
+        '[{id?: integer} | null]'
+    """
+    return _render(json.loads(shape))
+
+
+def _render(shape: typing.Any) -> str:
+    if isinstance(shape, str):
+        return shape
+    kind, body = next(iter(shape.items()))
+    if kind == "array":
+        return f"[{_render_value(body)}]"
+    if kind == "map":
+        return f"{{[key]: {_render_value(body)}}}"
+    fields = ", ".join(
+        f"{f['name']}{'?' if f['optional'] else ''}: {_render_value(f['value'])}" for f in body
+    )
+    return f"{{{fields}}}"
+
+
+def _render_value(value: Mapping[str, typing.Any]) -> str:
+    text = _render(value["shape"])
+    return f"{text} | null" if value["nullable"] else text
+
+
+def summarize(problems: list[str]) -> str:
+    """The first few problems on one line, the rest counted."""
+    shown = "; ".join(problems[:_MAX_PROBLEMS])
+    extra = len(problems) - _MAX_PROBLEMS
+    return f"{shown}; and {extra} more" if extra > 0 else shown
+
+
+def _head(node: ast.expr, ctx: _Context) -> str | None:
+    """The name a (possibly subscripted) annotation is, ``typing`` aliases resolved."""
+    if isinstance(node, ast.Subscript):
+        node = node.value
+    if isinstance(node, ast.Name):
+        return ctx.aliases.get(node.id, node.id)
+    if isinstance(node, ast.Attribute):
+        return node.attr
+    return None
+
+
+def _is_none(node: ast.expr) -> bool:
+    return isinstance(node, ast.Constant) and node.value is None
+
+
+def _members(node: ast.expr, ctx: _Context, depth: int = 0) -> list[ast.expr]:
+    """The alternatives of a union, ``Optional`` and ``Annotated`` looked through."""
+    if depth > _MAX_DEPTH:
+        return [node]
+    if isinstance(node, ast.Constant) and isinstance(node.value, str):
+        try:
+            return _members(ast.parse(node.value, mode="eval").body, ctx, depth + 1)
+        except SyntaxError, ValueError:
+            return [node]
+    if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
+        return _members(node.left, ctx, depth + 1) + _members(node.right, ctx, depth + 1)
+    if isinstance(node, ast.Subscript):
+        head = _head(node, ctx)
+        args = list(node.slice.elts) if isinstance(node.slice, ast.Tuple) else [node.slice]
+        if head == "Optional":
+            return [*_members(args[0], ctx, depth + 1), ast.Constant(None)]
+        if head == "Union":
+            return [m for arg in args for m in _members(arg, ctx, depth + 1)]
+        if head == "Annotated":
+            return _members(args[0], ctx, depth + 1)
+    return [node]
+
+
+def _admits_none(members: list[ast.expr], ctx: _Context) -> bool:
+    for member in members:
+        if _is_none(member) or _head(member, ctx) in _OPEN:
+            return True
+        if _head(member, ctx) == "Literal" and isinstance(member, ast.Subscript):
+            args = member.slice.elts if isinstance(member.slice, ast.Tuple) else [member.slice]
+            if any(_is_none(arg) for arg in args):
+                return True
+    return False
+
+
+def _describe(shape: typing.Any) -> str:
+    if isinstance(shape, str):
+        return _SCALAR_NAMES.get(shape, shape)
+    return {"array": "an array", "map": "a map", "object": "an object"}[next(iter(shape))]
+
+
+def _at(path: str, text: str) -> str:
+    return f"{path} {text}" if path else text
+
+
+def _check(
+    value: Mapping[str, typing.Any],
+    node: ast.expr,
+    path: str,
+    ctx: _Context,
+    out: list[str],
+    depth: int,
+) -> None:
+    if depth > _MAX_DEPTH:
+        return
+    members = _members(node, ctx)
+    if value["nullable"] and not _admits_none(members, ctx):
+        out.append(_at(path, f"can be null but is declared {ast.unparse(node)}"))
+    shape = value["shape"]
+    if shape in ("any", "never", "null"):
+        return
+    candidates = [m for m in members if not _is_none(m)]
+    if not candidates:
+        out.append(_at(path, f"is None but the query builds {_describe(shape)}"))
+        return
+    failures: list[list[str]] = []
+    for member in candidates:
+        found: list[str] = []
+        _check_member(shape, member, path, ctx, found, depth)
+        if not found:
+            return
+        failures.append(found)
+    out.extend(failures[0])
+
+
+def _mismatch(shape: typing.Any, member: ast.expr, path: str, out: list[str]) -> None:
+    out.append(_at(path, f"is {ast.unparse(member)} but the query builds {_describe(shape)}"))
+
+
+def _check_member(
+    shape: typing.Any,
+    member: ast.expr,
+    path: str,
+    ctx: _Context,
+    out: list[str],
+    depth: int,
+) -> None:
+    head = _head(member, ctx)
+    if head is None or head in _OPEN:
+        return
+    info = ctx.defs.get(head) if isinstance(member, ast.Name) else None
+    row_type = info if info is not None and info.kind in ("TypedDict", "BaseModel") else None
+    if isinstance(shape, str):
+        if head in _BUILTINS | _LISTS | _MAPPINGS or row_type:
+            if head not in _SCALAR_HEADS[shape]:
+                _mismatch(shape, member, path, out)
+        return
+    kind = next(iter(shape))
+    if kind == "array":
+        if head in _LISTS:
+            if isinstance(member, ast.Subscript):
+                _check(shape["array"], member.slice, f"{path}[]", ctx, out, depth + 1)
+        elif head in _BUILTINS | _MAPPINGS or row_type:
+            _mismatch(shape, member, path, out)
+    elif kind == "map":
+        if head in _MAPPINGS:
+            if parts := _mapping_parts(member, path, ctx, out):
+                _check(shape["map"], parts, f"{path}[key]", ctx, out, depth + 1)
+        elif head in _BUILTINS | _LISTS or row_type:
+            _mismatch(shape, member, path, out)
+    elif head in _MAPPINGS:
+        if parts := _mapping_parts(member, path, ctx, out):
+            for field in shape["object"]:
+                _check(field["value"], parts, _join(path, field["name"]), ctx, out, depth + 1)
+    elif row_type is not None:
+        _check_row_type(shape["object"], row_type, path, ctx, out, depth)
+    elif head in _BUILTINS | _LISTS:
+        _mismatch(shape, member, path, out)
+
+
+def _mapping_parts(member: ast.expr, path: str, ctx: _Context, out: list[str]) -> ast.expr | None:
+    """The value annotation of ``dict[K, V]``, reporting a key type that is not ``str``."""
+    if not isinstance(member, ast.Subscript) or not isinstance(member.slice, ast.Tuple):
+        return None
+    key, item = member.slice.elts[0], member.slice.elts[-1]
+    if (head := _head(key, ctx)) in _BUILTINS and head != "str":
+        out.append(_at(path, f"has keys of {ast.unparse(key)} but JSON object keys are strings"))
+    return item
+
+
+def _join(path: str, name: str) -> str:
+    return f"{path}.{name}" if path else name
+
+
+def _check_row_type(
+    fields: list[Mapping[str, typing.Any]],
+    info: TypeInfo,
+    path: str,
+    ctx: _Context,
+    out: list[str],
+    depth: int,
+) -> None:
+    inner = _Context(ctx.defs, info.aliases or ctx.aliases)
+    built = {field["name"]: field for field in fields}
+    # A Pydantic model reads keys by alias; its field types are not checked here.
+    model = info.kind == "BaseModel"
+    admits_extra = info.extra_items or info.allows_extra or info.opaque
+    for name, field in built.items():
+        where = _join(path, name)
+        key = info.field_for(name)
+        if key is None:
+            if not admits_extra:
+                out.append(
+                    _at(path or "the object", f"has key '{name}' that {info.name} does not declare")
+                )
+            continue
+        if field["optional"] and key in info.required:
+            out.append(_at(where, f"is required by {info.name} but the query may omit it"))
+        if not model and (annotation := info.annotations.get(key)) is not None:
+            _check(field["value"], annotation, where, inner, out, depth + 1)
+    for key, _ in info.fields:
+        names = info.column_names(key)
+        if key in info.required and key not in info.opaque and not built.keys() & set(names):
+            out.append(
+                _at(
+                    _join(path, names[0]),
+                    f"is required by {info.name} but the query does not build it",
+                )
+            )

sqlty/_json_types.py

@@ -0,0 +1,225 @@
+"""TypedDicts for JSON columns: what ``json-types = "typing.TypedDict"`` writes into a module.
+
+``rewrite`` types a JSON column ``typing.Any``. Where the engine knows the shape of the
+JSON a query builds (``jsonb_build_object()``, ``jsonb_agg()``, ``to_jsonb()``), the
+``typing.Any`` items of a ``SQL[tuple[...]]`` annotation become nested types instead, with a
+``TypedDict`` per object. Each class is written after the imports, below a
+``# sqlty: generated`` line, and ``strip`` removes them again.
+"""
+
+from __future__ import annotations
+
+import ast
+import json
+import keyword
+import re
+import typing
+
+from sqlty._core import analyze_columns
+from sqlty._source import _wrapper_name
+
+if typing.TYPE_CHECKING:
+    from sqlty._core import Schema
+
+MARKER = "# sqlty: generated"
+_BLOCK = (
+    rf"{MARKER}\r?\n(?:class [^\r\n]*:\r?\n(?:    [^\r\n]*\r?\n)+"
+    r"|[A-Za-z_]\w* = typing\.TypedDict\([^\r\n]*\)\r?\n)"
+)
+_GROUP = re.compile(rf"(\r?\n){{2}}{_BLOCK}(?:(?:\r?\n){{2}}{_BLOCK})*(?:\r?\n){{2}}")
+_SCALARS = {
+    "integer": "int",
+    "number": "float",
+    "string": "str",
+    "boolean": "bool",
+    "null": "None",
+    "any": "typing.Any",
+    "never": "typing.Any",
+}
+
+
+def add_json_types(
+    source: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+    import_style: str | None = None,
+) -> str:
+    """``source`` with the ``typing.Any`` row items of JSON columns of known shape typed.
+
+    Example:
+        >>> source = (
+        ...     "import typing\\nfrom sqlty import SQL\\n"
+        ...     "q = SQL[tuple[int, typing.Any]]"
+        ...     "(\\"--sql\\\\nSELECT 1::INTEGER AS id, jsonb_build_object('n', 1) AS doc\\")\\n"
+        ... )
+        >>> print(add_json_types(source))
+        import typing
+        from sqlty import SQL
+        <BLANKLINE>
+        <BLANKLINE>
+        # sqlty: generated
+        class Doc(typing.TypedDict):
+            n: int
+        <BLANKLINE>
+        <BLANKLINE>
+        q = SQL[tuple[int, Doc]]("--sql\\nSELECT 1::INTEGER AS id, jsonb_build_object('n', 1) AS doc")
+        <BLANKLINE>
+    """
+    wrapper = _wrapper_name(import_style)
+    try:
+        tree = ast.parse(source)
+    except SyntaxError, ValueError:
+        return source
+    if wrapper is None:
+        return source
+    module = _Module(tree, source)
+    edits: list[tuple[int, int, str]] = []
+    for node in ast.walk(tree):
+        if not (
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Subscript)
+            and ast.unparse(node.func.value) == wrapper
+            and node.args
+            and isinstance(node.args[0], ast.Constant)
+            and isinstance(node.args[0].value, str)
+        ):
+            continue
+        row = node.func.slice
+        if not (isinstance(row, ast.Subscript) and ast.unparse(row.value) == "tuple"):
+            continue
+        items = row.slice.elts if isinstance(row.slice, ast.Tuple) else [row.slice]
+        if not any(ast.unparse(item) in ("typing.Any", "Any") for item in items):
+            continue
+        try:
+            columns = analyze_columns(node.args[0].value, dialect=dialect, schema_sql=schema_sql)
+        except RuntimeError:
+            continue
+        if columns is None or len(columns) != len(items):
+            continue
+        for item, (name, _, _, nullable, shape) in zip(items, columns, strict=True):
+            if shape is None or ast.unparse(item) not in ("typing.Any", "Any"):
+                continue
+            value = {"shape": json.loads(shape), "nullable": nullable is True}
+            if value["shape"] == "any":
+                continue
+            edits.append((*module.span(item), module.type_of(value, _pascal(name))))
+    if not edits:
+        return source
+    for start, end, text in sorted(edits, reverse=True):
+        source = source[:start] + text + source[end:]
+    return module.with_classes(source)
+
+
+def strip_json_types(source: str) -> str:
+    """``source`` without the classes ``add_json_types`` wrote.
+
+    Example:
+        >>> strip_json_types("import typing\\n\\n\\n# sqlty: generated\\nclass A(typing.TypedDict):\\n    a: int\\n\\n\\nq = 1\\n")
+        'import typing\\nq = 1\\n'
+    """
+    return _GROUP.sub("", source) if MARKER in source else source
+
+
+class _Module:
+    """The top-level names of a module and the classes generated for it so far."""
+
+    def __init__(self, tree: ast.Module, source: str) -> None:
+        self.tree = tree
+        self.lines = source.splitlines(keepends=True)
+        self.taken: set[str] = set()
+        self.existing: dict[str, str] = {}
+        self.new: dict[str, str] = {}
+        generated = {
+            index + 1 for index, line in enumerate(self.lines) if line.rstrip("\r\n") == MARKER
+        }
+        for stmt in tree.body:
+            names = [
+                target.id for target in getattr(stmt, "targets", []) if isinstance(target, ast.Name)
+            ]
+            if isinstance(stmt, ast.ClassDef | ast.FunctionDef | ast.AsyncFunctionDef):
+                names.append(stmt.name)
+            if isinstance(stmt, ast.Import | ast.ImportFrom):
+                names += [(alias.asname or alias.name).split(".")[0] for alias in stmt.names]
+            self.taken.update(names)
+            if stmt.lineno - 1 in generated and names:
+                self.existing[names[0]] = ast.unparse(stmt)
+
+    def span(self, node: ast.expr) -> tuple[int, int]:
+        """Character offsets of ``node`` in the source (``ast`` counts UTF-8 bytes)."""
+        assert node.end_lineno is not None and node.end_col_offset is not None
+
+        def offset(line: int, column: int) -> int:
+            before = sum(len(text) for text in self.lines[: line - 1])
+            return before + len(self.lines[line - 1].encode()[:column].decode())
+
+        return offset(node.lineno, node.col_offset), offset(node.end_lineno, node.end_col_offset)
+
+    def type_of(self, value: dict[str, typing.Any], name: str) -> str:
+        """The Python type of a JSON value of the engine's shape; objects get a class ``name``."""
+        shape = value["shape"]
+        if isinstance(shape, str):
+            text = _SCALARS.get(shape, "typing.Any")
+        else:
+            kind, body = next(iter(shape.items()))
+            if kind == "array":
+                text = f"list[{self.type_of(body, name + 'Item')}]"
+            elif kind == "map":
+                text = f"dict[str, {self.type_of(body, name + 'Value')}]"
+            else:
+                text = self.typed_dict(body, name)
+        return f"{text} | None" if value["nullable"] and text != "None" else text
+
+    def typed_dict(self, fields: list[dict[str, typing.Any]], name: str) -> str:
+        """The name of a ``TypedDict`` class with ``fields``, written once per module."""
+        types = []
+        for field in fields:
+            text = self.type_of(field["value"], name + _pascal(field["name"]))
+            types.append(
+                (field["name"], f"typing.NotRequired[{text}]" if field["optional"] else text)
+            )
+        for candidate in [name, *(f"{name}{n}" for n in range(2, 100))]:
+            code = _class(candidate, types)
+            body = ast.unparse(ast.parse(code).body[0])
+            if self.existing.get(candidate) == body or self.new.get(candidate) == code:
+                return candidate
+            if candidate not in self.taken and candidate not in self.new:
+                self.new[candidate] = code
+                return candidate
+        return "typing.Any"
+
+    def with_classes(self, source: str) -> str:
+        """``source`` with the new classes after the imports, or after the generated ones."""
+        if not self.new:
+            return source
+        eol = "\r\n" if "\r\n" in source else "\n"
+        blocks = "\n\n".join(f"{MARKER}\n{code}" for code in self.new.values())
+        blocks = blocks.replace("\n", eol)
+        groups = list(_GROUP.finditer(source))
+        if groups:
+            at = groups[-1].end()
+            return source[:at] + blocks + eol * 2 + source[at:]
+        header = [stmt for stmt in self.tree.body if isinstance(stmt, ast.Import | ast.ImportFrom)]
+        line = header[-1].end_lineno if header else 0
+        at = sum(len(text) for text in self.lines[: line or 0])
+        if at and not source[:at].endswith("\n"):
+            source = source[:at] + eol + source[at:]
+            at += len(eol)
+        return source[:at] + eol * 2 + blocks + eol * 2 + source[at:]
+
+
+def _class(name: str, fields: list[tuple[str, str]]) -> str:
+    if fields and all(key.isidentifier() and not keyword.iskeyword(key) for key, _ in fields):
+        body = "".join(f"    {key}: {text}\n" for key, text in fields)
+        return f"class {name}(typing.TypedDict):\n{body}"
+    items = ", ".join(f"{key!r}: {text}" for key, text in fields)
+    return f"{name} = typing.TypedDict({name!r}, {{{items}}})\n"
+
+
+def _pascal(name: str) -> str:
+    """``order_lines`` as ``OrderLines``; a name with no letters is ``Json``."""
+    words = re.findall(r"[A-Za-z0-9]+", name)
+    text = "".join(word[:1].upper() + word[1:] for word in words)
+    if not text or text[0].isdigit():
+        text = "Json" + text
+    return text