sqlty 0.1.0__cp315-cp315-win_amd64.whl

sqlty/_core.pyi

@@ -0,0 +1,108 @@
+"""Type stubs for the Rust extension module (sqlty._core)."""
+
+from typing import final
+
+@final
+class Schema:
+    @property
+    def dialect(self) -> str: ...
+
+def load_schema(
+    schema_sql: str, *, dialect: str = "postgres", types: list[tuple[str, str]] | None = None
+) -> Schema: ...
+def infer_type(
+    sql: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+) -> str | None: ...
+def analyze_select(
+    sql: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+) -> list[tuple[str, str | None, str | None]] | None: ...
+def explain_select(
+    sql: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+) -> list[tuple[str, tuple[str, str, str] | None, tuple[str, str, str] | None]] | None: ...
+def analyze_columns(
+    sql: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+) -> list[tuple[str, str | None, str | None, bool | None, str | None]] | None: ...
+def analyze_params(
+    sql: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+) -> list[tuple[int, str | None, str | None, str | None, bool | None, bool]] | None: ...
+def composite_fields(
+    type_name: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+) -> list[tuple[str, str | None, str | None, bool | None]] | None: ...
+def query_diagnostics(
+    sql: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+) -> list[tuple[str, str]]: ...
+def name_rules() -> list[str]: ...
+def dialect_names() -> list[str]: ...
+def python_types() -> list[str]: ...
+def parse_error(sql: str, *, dialect: str = "postgres") -> str | None: ...
+def returns_rows(sql: str, *, dialect: str = "postgres") -> bool: ...
+def rewrite_source(
+    source: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+    import_style: str | None = None,
+    untyped: bool = False,
+    keep_lines: bool = False,
+) -> str: ...
+def strip_source(source: str, *, import_style: str | None = None) -> str: ...
+def fix_source(
+    source: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+    import_style: str | None = None,
+    untyped: bool = False,
+    rules: dict[str, str] | None = None,
+    suppressed: dict[int, list[str] | None] | None = None,
+) -> str: ...
+def expand_star(
+    sql: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+) -> str | None: ...
+def select_star_finding(
+    sql: str, *, dialect: str = "postgres", has_schema: bool = False
+) -> tuple[str, str] | None: ...
+def check_source(
+    source: str,
+    *,
+    dialect: str = "postgres",
+    schema_sql: str | Schema | None = None,
+    rules: dict[str, str] | None = None,
+    import_style: str | None = None,
+    untyped: bool = False,
+    any_calls: list[int] = ...,
+) -> list[dict[str, str | int | None]]: ...
+def wrapped_queries(
+    source: str, *, import_style: str | None = None
+) -> list[tuple[int, str, str | None]]: ...
+def query_sites(
+    source: str, *, import_style: str | None = None, any_calls: list[int] = ...
+) -> list[tuple[int, str, str]]: ...
+def schema_skipped_tables(schema_sql: str | Schema, *, dialect: str = "postgres") -> list[str]: ...
+def schema_skipped_objects(
+    schema_sql: str | Schema, *, dialect: str = "postgres"
+) -> list[tuple[str, str]]: ...

sqlty/_decode.py

@@ -0,0 +1,62 @@
+r"""The ``sqlty`` codec's decoder, without registering the codec.
+
+``decode`` turns source bytes into text with ``--sql`` strings wrapped as
+``SQL[T](...)``, as ``sqlty rewrite`` does, and adds the imports they need:
+
+    >>> from sqlty._decode import decode
+    >>> decode(b'q = "--sql\\nSELECT 1::INTEGER"')[0]
+    'from sqlty import SQL; q = SQL[tuple[int]]("--sql\\nSELECT 1::INTEGER")'
+
+Strings whose row type can't be inferred (invalid SQL, DML without
+RETURNING, DDL, empty queries) are left as they are. The imports are joined
+with ``;`` onto an existing line (after the docstring or ``__future__``
+imports, else before the first statement) rather than a new line above
+the code: when CPython runs a file as a script, it decodes from the end of the
+coding line and drops the first decoded line (#111). Line numbers match the file
+on disk, except when a compound statement (``def``, ``class``, ``@``, ``if``,
+...) comes right after the coding line, with no docstring, blank or comment
+line in between; then one line is added. The ``# coding: sqlty``
+declaration becomes ``utf-8``, so the decoded text runs without the codec.
+``import-style`` comes from ``[tool.sqlty]`` in pyproject.toml; on any error the
+source is returned unchanged. Import ``sqlty.magic`` to register the codec.
+"""
+
+from __future__ import annotations
+
+import codecs
+import encodings
+import re
+from collections.abc import Buffer
+from functools import cache
+
+__all__ = ["decode"]
+
+_utf_8_codec = encodings.search_function("utf8")
+assert _utf_8_codec is not None, "UTF-8 codec not found"
+utf_8_codec: codecs.CodecInfo = _utf_8_codec
+
+# PEP 263 declaration of this codec on the first or second line.
+_CODING = re.compile(r"\A((?:[^\n]*\n)?[ \t\f]*#[^\n]*?coding[:=][ \t]*)sqlty\b")
+
+
+@cache
+def _import_style() -> str | None:
+    try:
+        from sqlty._config import load_config
+
+        config, _ = load_config()
+        return config.get("import_style")
+    except Exception:
+        return None
+
+
+def decode(b: Buffer, errors: str = "strict") -> tuple[str, int]:
+    """Decode UTF-8 ``b``, wrapping ``--sql`` strings; returns ``(text, length)``."""
+    source, length = utf_8_codec.decode(b, errors)
+    try:
+        from sqlty.core import rewrite_source
+
+        decoded = rewrite_source(source, import_style=_import_style(), keep_lines=True)
+        return _CODING.sub(r"\1utf-8", decoded, count=1), length
+    except Exception:
+        return source, length

sqlty/_drivers.py

@@ -0,0 +1,274 @@
+"""Driver calls that pass a query and its parameters: ``cursor.execute(query, params)``.
+
+A receiver is recognised only when the file shows what it is: a variable bound to
+``psycopg.connect(...)``, ``sqlite3.connect(...)`` or ``asyncpg.connect(...)`` (also as
+``with ... as name``), to ``<connection>.cursor()``, or annotated with a driver class. The
+query is a string literal, ``SQL[T]("...")`` of one, or a name assigned exactly one of
+those. Anything else is skipped.
+
+Names are resolved per scope, and a name a loop, comprehension, unpacking, import or
+unannotated parameter binds is unknown.
+
+Example:
+    >>> src = '''
+    ... import sqlite3
+    ... conn = sqlite3.connect("db")
+    ... conn.execute("SELECT 1 WHERE ? = 1", (1, 2))
+    ... '''
+    >>> [(c.line, c.query, [a.key for a in c.params or ()]) for c in find_driver_calls(src)]
+    [(4, 'SELECT 1 WHERE ? = 1', [1, 2])]
+"""
+
+from __future__ import annotations
+
+import ast
+from collections import defaultdict
+
+from sqlty._helpers import HelperCall
+from sqlty._params import ParamArg, arg_type, literal_params
+from sqlty._source import ParsedSource
+
+# Driver factories and classes, by the dotted name they are imported as, to the driver.
+_CONNECTIONS = {
+    "psycopg.connect": "psycopg",
+    "psycopg.Connection": "psycopg",
+    "psycopg.AsyncConnection": "psycopg",
+    "psycopg.AsyncConnection.connect": "psycopg",
+    "psycopg.Connection.connect": "psycopg",
+    "sqlite3.connect": "sqlite3",
+    "sqlite3.Connection": "sqlite3",
+    "asyncpg.connect": "asyncpg",
+    "asyncpg.Connection": "asyncpg",
+    "asyncpg.create_pool": "asyncpg",
+    "asyncpg.Pool": "asyncpg",
+}
+_CURSORS = {
+    "psycopg.Cursor": "psycopg",
+    "psycopg.AsyncCursor": "psycopg",
+    "sqlite3.Cursor": "sqlite3",
+}
+_METHODS = {
+    "psycopg": {"execute", "executemany"},
+    "sqlite3": {"execute", "executemany"},
+    "asyncpg": {"execute", "fetch", "fetchrow", "fetchval", "executemany"},
+}
+_PARAM_KEYWORD = {"psycopg": "params", "sqlite3": "parameters"}
+
+
+def _dotted(node: ast.expr) -> str | None:
+    parts: list[str] = []
+    while isinstance(node, ast.Attribute):
+        parts.append(node.attr)
+        node = node.value
+    return ".".join([node.id, *reversed(parts)]) if isinstance(node, ast.Name) else None
+
+
+def _qualified(name: str | None, parsed: ParsedSource) -> str | None:
+    """``name`` with its first part replaced by what it was imported as."""
+    if name is None:
+        return None
+    head, _, rest = name.partition(".")
+    imported = parsed.imported_names.get(head)
+    return None if imported is None else f"{imported}.{rest}" if rest else imported
+
+
+_Binding = tuple[ast.expr | None, ast.expr | None]  # (value, annotation); both None: unknown
+_Scope = ast.AST  # the module, or a function
+
+
+class _Scopes:
+    """What names mean in each scope of a file.
+
+    A name takes its meaning from the innermost enclosing function (or the module) that
+    binds it. Names bound by anything but a plain assignment, an annotated parameter or
+    ``with ... as name`` (loops, comprehensions, unpacking, imports, definitions, an
+    unannotated parameter) are unknown.
+    """
+
+    def __init__(self, parsed: ParsedSource) -> None:
+        assert parsed.tree is not None
+        self.parsed = parsed
+        self.tree: _Scope = parsed.tree
+        self._bindings: dict[int, dict[str, list[_Binding]]] = {}
+        self._resolving: set[tuple[int, str]] = set()
+
+    def chain(self, call: ast.Call) -> list[_Scope]:
+        """The scopes around ``call``, outermost (the module) first."""
+        inner = [
+            s
+            for s in self.parsed.scopes_of(call)
+            if isinstance(s, (ast.FunctionDef, ast.AsyncFunctionDef))
+        ]
+        return [self.tree, *inner]
+
+    def _of(self, scope: _Scope) -> dict[str, list[_Binding]]:
+        if (found := self._bindings.get(id(scope))) is not None:
+            return found
+        names: dict[str, list[_Binding]] = defaultdict(list)
+        self._bindings[id(scope)] = names
+        unknown: _Binding = (None, None)
+
+        def bind(target: ast.AST, binding: _Binding = unknown) -> None:
+            if isinstance(target, ast.Name):
+                names[target.id].append(binding)
+            else:
+                for node in ast.walk(target):
+                    if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store):
+                        names[node.id].append(unknown)
+
+        if isinstance(scope, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            args = scope.args
+            for arg in [*args.posonlyargs, *args.args, *args.kwonlyargs, args.vararg, args.kwarg]:
+                if arg is not None:
+                    names[arg.arg].append((None, arg.annotation))
+        pending = list(ast.iter_child_nodes(scope))
+        while pending:
+            node = pending.pop()
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
+                names[node.name].append(unknown)
+                continue
+            if isinstance(node, ast.Lambda):
+                continue
+            if isinstance(node, ast.Assign):
+                for target in node.targets:
+                    bind(target, (node.value, None) if isinstance(target, ast.Name) else unknown)
+            elif isinstance(node, ast.AnnAssign):
+                bind(node.target, (node.value, node.annotation))
+            elif isinstance(node, (ast.With, ast.AsyncWith)):
+                for item in node.items:
+                    if item.optional_vars is not None:
+                        bind(item.optional_vars, (item.context_expr, None))
+            elif isinstance(
+                node, (ast.For, ast.AsyncFor, ast.comprehension, ast.AugAssign, ast.NamedExpr)
+            ):
+                bind(node.target)
+            elif isinstance(node, ast.ExceptHandler) and node.name:
+                names[node.name].append(unknown)
+            elif isinstance(node, (ast.Import, ast.ImportFrom)):
+                for alias in node.names:
+                    names[(alias.asname or alias.name).partition(".")[0]].append(unknown)
+            pending.extend(ast.iter_child_nodes(node))
+        return names
+
+    def _lookup(self, name: str, scopes: list[_Scope]) -> list[_Binding] | None:
+        for scope in reversed(scopes):
+            if (found := self._of(scope).get(name)) is not None:
+                return found
+        return None
+
+    def driver(self, name: str, scopes: list[_Scope]) -> str | None:
+        """The driver of the connection or cursor ``name`` is, in ``scopes``, when every binding agrees."""
+        key = (id(scopes[-1]), name)
+        if key in self._resolving or (bindings := self._lookup(name, scopes)) is None:
+            return None
+        self._resolving.add(key)
+        try:
+            drivers = {self._kind(value, annotation, scopes) for value, annotation in bindings}
+        finally:
+            self._resolving.discard(key)
+        return drivers.pop() if len(drivers) == 1 else None
+
+    def _kind(
+        self, value: ast.expr | None, annotation: ast.expr | None, scopes: list[_Scope]
+    ) -> str | None:
+        if annotation is not None:
+            name = _qualified(_dotted(annotation), self.parsed)
+            if name in _CONNECTIONS:
+                return _CONNECTIONS[name]
+            if name in _CURSORS:
+                return _CURSORS[name]
+        while isinstance(value, ast.Await):
+            value = value.value
+        if isinstance(value, ast.Call):
+            name = _qualified(_dotted(value.func), self.parsed)
+            if name in _CONNECTIONS:
+                return _CONNECTIONS[name]
+            func = value.func
+            if (
+                isinstance(func, ast.Attribute)
+                and func.attr == "cursor"
+                and isinstance(func.value, ast.Name)
+            ):
+                return self.driver(func.value.id, scopes)
+        return None
+
+    def query(self, name: str, scopes: list[_Scope]) -> str | None:
+        """The query literal ``name`` is bound to, in ``scopes``, when it has exactly one binding."""
+        bindings = self._lookup(name, scopes)
+        if bindings is None or len(bindings) != 1:
+            return None
+        return _literal(bindings[0][0])
+
+
+def _literal(node: ast.expr | None) -> str | None:
+    """The text of a string literal, possibly wrapped as ``SQL[T]("...")``."""
+    if isinstance(node, ast.Call) and isinstance(node.func, ast.Subscript) and len(node.args) == 1:
+        if _dotted(node.func.value) in {"SQL", "sqlty.SQL"}:
+            return _literal(node.args[0])
+    if isinstance(node, ast.Constant) and isinstance(node.value, str):
+        return node.value
+    return None
+
+
+def find_driver_calls(source: str | ParsedSource) -> list[HelperCall]:
+    """Calls of ``execute``-like methods on a recognised connection or cursor, with their query and parameters.
+
+    A call with no parameters has ``params == ()``, so a query that has placeholders is
+    reported. ``executemany`` gives one entry per row of a list or tuple literal.
+    """
+    parsed = ParsedSource.of(source)
+    if not any(driver in parsed.text for driver in ("psycopg", "sqlite3", "asyncpg")):
+        return []
+    if parsed.tree is None:
+        return []
+    scopes = _Scopes(parsed)
+    found: list[HelperCall] = []
+    aliases = parsed.typing_aliases
+    for call in parsed.calls_to(None):
+        func = call.func
+        if not (
+            isinstance(func, ast.Attribute)
+            and isinstance(func.value, ast.Name)
+            and (driver := scopes.driver(func.value.id, chain := scopes.chain(call)))
+            and func.attr in _METHODS[driver]
+            and call.args
+            and not any(isinstance(arg, ast.Starred) for arg in call.args)
+        ):
+            continue
+        first = call.args[0]
+        query = _literal(first)
+        if query is None and isinstance(first, ast.Name):
+            query = scopes.query(first.id, chain)
+        if query is None:
+            continue
+        stack = parsed.scopes_of(call)
+        if driver == "asyncpg":
+            if func.attr == "executemany":
+                continue
+            params: tuple[ParamArg, ...] | None = tuple(
+                ParamArg(i, arg_type(arg, parsed, stack)) for i, arg in enumerate(call.args[1:], 1)
+            )
+            found.append(HelperCall(call.lineno, query, None, aliases, params))
+            continue
+        keyword = _PARAM_KEYWORD[driver]
+        arg = call.args[1] if len(call.args) > 1 else None
+        if arg is None:
+            arg = next((k.value for k in call.keywords if k.arg == keyword), None)
+            if arg is None and any(k.arg is None for k in call.keywords):
+                continue
+        if func.attr == "executemany":
+            if isinstance(arg, (ast.List, ast.Tuple)):
+                for row in arg.elts:
+                    if (params := literal_params(row, parsed, stack)) is not None:
+                        found.append(HelperCall(call.lineno, query, None, aliases, params))
+            continue
+        found.append(
+            HelperCall(
+                call.lineno,
+                query,
+                None,
+                aliases,
+                () if arg is None else literal_params(arg, parsed, stack),
+            )
+        )
+    return found

sqlty/_gitignore.py

@@ -0,0 +1,58 @@
+"""Which paths of a git work tree its ``.gitignore`` files ignore."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from pathspec import GitIgnoreSpec
+
+_Chain = tuple[tuple[Path, GitIgnoreSpec], ...]
+
+
+class GitIgnore:
+    """The ignore rules of the repositories a walk enters, read once per directory.
+
+    A directory's rules are those of ``.git/info/exclude`` and every ``.gitignore``
+    from the repository root down to it; a deeper file overrides a shallower one,
+    and a later line an earlier one, as in git. Outside a repository nothing is
+    ignored.
+
+    Example:
+        >>> GitIgnore().ignored(Path("/"), "tmp", is_dir=True)
+        False
+    """
+
+    def __init__(self) -> None:
+        self._chains: dict[Path, _Chain | None] = {}
+
+    def ignored(self, directory: Path, name: str, *, is_dir: bool) -> bool:
+        """Whether the entry ``name`` of the resolved ``directory`` is ignored."""
+        verdict = False
+        for base, spec in self._chain(directory) or ():
+            relative = (directory / name).relative_to(base).as_posix()
+            if (
+                include := spec.check_file(relative + "/" if is_dir else relative).include
+            ) is not None:
+                verdict = include
+        return verdict
+
+    def _chain(self, directory: Path) -> _Chain | None:
+        if directory not in self._chains:
+            if (directory / ".git").exists():
+                chain: _Chain | None = _read(directory, directory / ".git" / "info" / "exclude")
+            elif directory.parent == directory:
+                chain = None
+            else:
+                chain = self._chain(directory.parent)
+            if chain is not None:
+                chain += _read(directory, directory / ".gitignore")
+            self._chains[directory] = chain
+        return self._chains[directory]
+
+
+def _read(base: Path, path: Path) -> _Chain:
+    try:
+        lines = path.read_text(encoding="utf-8").splitlines()
+    except OSError, UnicodeDecodeError:
+        return ()
+    return ((base, GitIgnoreSpec.from_lines(lines)),)