remote/sqlite-utils
1
0
Fork 0
mirror of https://github.com/simonw/sqlite-utils.git synced 2026-09-08 01:24:20 +02:00
Code Issues Projects Releases Packages Wiki Activity

Merge remote-tracking branch 'origin/main' into fix/convert-dryrun-quote-identifier

Browse source 
# Conflicts:
#	sqlite_utils/cli.py
This commit is contained in:
Simon Willison 2026-08-12 14:07:24 -07:00
commit fc2dcfca36
71 changed files with 4489 additions and 2191 deletions
Download patch file Download diff file Expand all files Collapse all files

7
sqlite_utils/__init__.py
View file

@ -1,7 +1,6 @@
from .utils import suggest_column_types
from .hookspecs import hookimpl
from .hookspecs import hookspec
from .db import Database
from .hookspecs import hookimpl, hookspec
from .migrations import Migrations
from .utils import suggest_column_types
__all__ = ["Database", "Migrations", "suggest_column_types", "hookimpl", "hookspec"]
__all__ = ["Database", "Migrations", "hookimpl", "hookspec", "suggest_column_types"]

19
sqlite_utils/cli.py
View file

@ -1093,7 +1093,7 @@ def insert_upsert_implementation(
column_type_overrides = {column: ctype.upper() for column, ctype in (types or [])}
def _insert_docs(docs, tracker=None):
extra_kwargs = {
extra_kwargs: dict[str, Any] = {
"ignore": ignore,
"replace": replace,
"truncate": truncate,
@ -2478,6 +2478,8 @@ def rows(
if limit:
sql += f" limit {limit}"
if offset:
if not limit:
sql += " limit -1"
sql += f" offset {offset}"
ctx.invoke(
query,
@ -3273,15 +3275,10 @@ def convert(
raise click.ClickException(str(e))
if dry_run:
# Pull first 20 values for first column and preview them
if multi:
def preview(v):
def preview(v):
if multi:
return json.dumps(fn(v), default=repr, ensure_ascii=False) if v else v
else:
def preview(v):
return fn(v) if v else v
return fn(v) if v else v
db.conn.create_function("preview_transform", 1, preview)
col_quoted = quote_identifier(columns[0])
@ -3788,12 +3785,12 @@ def _rows_from_code(code):
code = pathlib.Path(code).read_text()
except FileNotFoundError:
raise click.ClickException(f"File not found: {code}")
namespace = {}
namespace: dict[str, Any] = {}
try:
exec(code, namespace) # noqa: S102
except SyntaxError as ex:
raise click.ClickException(f"Error in --code: {ex}")
rows = namespace.get("rows")
rows: Any = namespace.get("rows")
if callable(rows):
rows = rows()
if isinstance(rows, dict):

676
sqlite_utils/create_table_parser.py Normal file
View file

@ -0,0 +1,676 @@
"""Helpers for parsing CHECK constraints from SQLite CREATE TABLE SQL.
SQLite does not expose CHECK constraints through a pragma, so preserving them
across a table rebuild requires reading ``sqlite_schema.sql``. This module is
deliberately small, but it uses a real lexer: strings, quoted identifiers and
comments are opaque, every token retains its source span and malformed input is
reported instead of being silently under-parsed.
"""
import re
from dataclasses import dataclass, field
from typing import Any
@dataclass
class Check:
check: str
name: str = ""
column: str = ""
options: list[Any] | None = None
# Source details are excluded from equality and repr so callers can compare
# semantic constraints while still having the original SQL available for
# diagnostics or future lossless edits.
sql: str = field(default="", compare=False, repr=False)
start: int = field(default=-1, compare=False, repr=False)
end: int = field(default=-1, compare=False, repr=False)
@dataclass(frozen=True)
class ColumnComments:
before: str = ""
after: str = ""
class ParseError(ValueError):
pass
@dataclass(frozen=True)
class _Token:
kind: str
text: str
start: int
end: int
def is_keyword(self, keyword: str) -> bool:
return self.kind == "word" and self.text.upper() == keyword
_PUNCTUATION = frozenset("(),.;+-*/%<>=!~|&?:")
_TRIVIA = frozenset(("whitespace", "comment"))
_TABLE_CONSTRAINT_KEYWORDS = frozenset(("PRIMARY", "UNIQUE", "CHECK", "FOREIGN"))
_OTHER_COLUMN_CONSTRAINT_KEYWORDS = frozenset(
("PRIMARY", "UNIQUE", "REFERENCES", "DEFAULT", "NOT", "COLLATE", "GENERATED")
)
_SQLITE_KEYWORDS = frozenset(
(
"ABORT",
"ACTION",
"ADD",
"AFTER",
"ALL",
"ALTER",
"ANALYZE",
"AND",
"AS",
"ASC",
"ATTACH",
"AUTOINCREMENT",
"BEFORE",
"BEGIN",
"BETWEEN",
"BY",
"CASCADE",
"CASE",
"CAST",
"CHECK",
"COLLATE",
"COLUMN",
"COMMIT",
"CONFLICT",
"CONSTRAINT",
"CREATE",
"CROSS",
"CURRENT_DATE",
"CURRENT_TIME",
"CURRENT_TIMESTAMP",
"DATABASE",
"DEFAULT",
"DEFERRABLE",
"DEFERRED",
"DELETE",
"DESC",
"DETACH",
"DISTINCT",
"DO",
"DROP",
"EACH",
"ELSE",
"END",
"ESCAPE",
"EXCEPT",
"EXCLUDE",
"EXCLUSIVE",
"EXISTS",
"EXPLAIN",
"FAIL",
"FALSE",
"FILTER",
"FIRST",
"FOLLOWING",
"FOR",
"FOREIGN",
"FROM",
"FULL",
"GENERATED",
"GLOB",
"GROUP",
"GROUPS",
"HAVING",
"IF",
"IGNORE",
"IMMEDIATE",
"IN",
"INDEX",
"INDEXED",
"INITIALLY",
"INNER",
"INSERT",
"INSTEAD",
"INTERSECT",
"INTO",
"IS",
"ISNULL",
"JOIN",
"KEY",
"LAST",
"LEFT",
"LIKE",
"LIMIT",
"MATCH",
"MATERIALIZED",
"NATURAL",
"NO",
"NOT",
"NOTHING",
"NOTNULL",
"NULL",
"NULLS",
"OF",
"OFFSET",
"ON",
"OR",
"ORDER",
"OTHERS",
"OUTER",
"OVER",
"PARTITION",
"PLAN",
"PRAGMA",
"PRECEDING",
"PRIMARY",
"QUERY",
"RAISE",
"RANGE",
"RECURSIVE",
"REFERENCES",
"REGEXP",
"REINDEX",
"RELEASE",
"RENAME",
"REPLACE",
"RESTRICT",
"RETURNING",
"RIGHT",
"ROLLBACK",
"ROW",
"ROWS",
"SAVEPOINT",
"SELECT",
"SET",
"STRICT",
"TABLE",
"TEMP",
"TEMPORARY",
"THEN",
"TIES",
"TO",
"TRANSACTION",
"TRIGGER",
"TRUE",
"UNBOUNDED",
"UNION",
"UNIQUE",
"UPDATE",
"USING",
"VACUUM",
"VALUES",
"VIEW",
"VIRTUAL",
"WHEN",
"WHERE",
"WINDOW",
"WITH",
"WITHOUT",
)
)
_INTEGER_RE = re.compile(r"[+-]?(?:0[xX][0-9a-fA-F]+|[0-9]+)\Z")
_FLOAT_RE = re.compile(
r"[+-]?(?:(?:[0-9]+\.[0-9]*|\.[0-9]+)(?:[eE][+-]?[0-9]+)?|"
r"[0-9]+[eE][+-]?[0-9]+)\Z"
)
def _lex(sql: str) -> list[_Token]:
tokens: list[_Token] = []
i = 0
while i < len(sql):
start = i
char = sql[i]
if char.isspace():
i += 1
while i < len(sql) and sql[i].isspace():
i += 1
tokens.append(_Token("whitespace", sql[start:i], start, i))
continue
if sql.startswith("--", i):
newline = sql.find("\n", i + 2)
i = len(sql) if newline == -1 else newline + 1
tokens.append(_Token("comment", sql[start:i], start, i))
continue
if sql.startswith("/*", i):
end = sql.find("*/", i + 2)
if end == -1:
raise ParseError("Unterminated SQL comment")
i = end + 2
tokens.append(_Token("comment", sql[start:i], start, i))
continue
if char in ("'", '"', "`"):
quote = char
i += 1
while i < len(sql):
if sql[i] == quote:
if i + 1 < len(sql) and sql[i + 1] == quote:
i += 2
continue
i += 1
break
i += 1
else:
raise ParseError(f"Unterminated {quote} quoted token")
kind = "string" if quote == "'" else "identifier"
tokens.append(_Token(kind, sql[start:i], start, i))
continue
if char == "[":
end = sql.find("]", i + 1)
if end == -1:
raise ParseError("Unterminated [ quoted identifier")
i = end + 1
tokens.append(_Token("identifier", sql[start:i], start, i))
continue
if char in _PUNCTUATION:
i += 1
tokens.append(_Token("punct", char, start, i))
continue
# SQLite accepts any character >= U+0080 in a bare identifier. More
# generally, consume until a lexical delimiter rather than relying on
# Python's narrower definition of an alphanumeric character.
i += 1
while i < len(sql):
if sql[i].isspace() or sql[i] in _PUNCTUATION or sql[i] in "'\"`[":
break
i += 1
tokens.append(_Token("word", sql[start:i], start, i))
return tokens
def _meaningful(tokens: list[_Token]) -> list[_Token]:
return [token for token in tokens if token.kind not in _TRIVIA]
def _unquote(token: str) -> str:
if len(token) >= 2 and token[0] in ("'", '"', "`") and token[-1] == token[0]:
return token[1:-1].replace(token[0] * 2, token[0])
if len(token) >= 2 and token[0] == "[" and token[-1] == "]":
return token[1:-1]
return token
def _matching_paren(tokens: list[_Token], open_index: int) -> int:
if tokens[open_index].text != "(":
raise ParseError("Expected an opening parenthesis")
depth = 0
for index in range(open_index, len(tokens)):
if tokens[index].text == "(":
depth += 1
elif tokens[index].text == ")":
depth -= 1
if depth == 0:
return index
raise ParseError("Unbalanced parentheses")
def _split_spans(sql: str, tokens: list[_Token]) -> list[tuple[str, int, int]]:
if not tokens:
return []
items: list[tuple[str, int, int]] = []
depth = 0
start = tokens[0].start
for token in tokens:
if token.text == "(":
depth += 1
elif token.text == ")":
depth -= 1
if depth < 0:
raise ParseError("Unbalanced parentheses")
elif token.text == "," and depth == 0:
raw = sql[start : token.start]
item = raw.strip()
if item:
item_start = start + len(raw) - len(raw.lstrip())
items.append((item, item_start, item_start + len(item)))
start = token.end
if depth:
raise ParseError("Unbalanced parentheses")
raw = sql[start : tokens[-1].end]
item = raw.strip()
if item:
item_start = start + len(raw) - len(raw.lstrip())
items.append((item, item_start, item_start + len(item)))
return items
def _split_ranges(sql: str, tokens: list[_Token]) -> list[str]:
return [item for item, _, _ in _split_spans(sql, tokens)]
def _strip_outer_parens(tokens: list[_Token]) -> list[_Token]:
while tokens and tokens[0].text == "(":
close = _matching_paren(tokens, 0)
if close != len(tokens) - 1:
break
tokens = tokens[1:-1]
return tokens
_NO_LITERAL = object()
def _literal_value(text: str) -> Any:
tokens = _meaningful(_lex(text))
if len(tokens) == 1 and tokens[0].kind == "string":
return _unquote(tokens[0].text)
raw = "".join(token.text for token in tokens)
if raw.upper() == "NULL":
return None
if raw.upper() == "TRUE":
return True
if raw.upper() == "FALSE":
return False
if _INTEGER_RE.fullmatch(raw):
try:
return (
int(raw, 16) if raw.lower().lstrip("+-").startswith("0x") else int(raw)
)
except ValueError:
return _NO_LITERAL
if _FLOAT_RE.fullmatch(raw):
try:
return float(raw)
except ValueError:
return _NO_LITERAL
return _NO_LITERAL
def _ascii_fold(identifier: str) -> str:
return identifier.translate(
str.maketrans("ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz")
)
def _parse_options(expression: str, column: str) -> list[Any] | None:
tokens = _strip_outer_parens(_meaningful(_lex(expression)))
if len(tokens) < 4:
return None
lhs = tokens[0]
if lhs.kind not in ("word", "identifier"):
return None
if column and _ascii_fold(_unquote(lhs.text)) != _ascii_fold(column):
return None
if not tokens[1].is_keyword("IN") or tokens[2].text != "(":
return None
close = _matching_paren(tokens, 2)
if close != len(tokens) - 1:
return None
inner = expression[tokens[2].end : tokens[close].start]
inner_tokens = _lex(inner)
if not _meaningful(inner_tokens):
return []
values = []
for item in _split_ranges(inner, inner_tokens):
value = _literal_value(item)
if value is _NO_LITERAL:
return None
values.append(value)
return values
def _check_after(
item: str,
tokens: list[_Token],
check_index: int,
name: str,
column: str,
constraint_start: int,
base_offset: int,
) -> tuple[Check, int]:
if check_index + 1 >= len(tokens) or tokens[check_index + 1].text != "(":
raise ParseError("CHECK must be followed by a parenthesized expression")
close = _matching_paren(tokens, check_index + 1)
expression = item[tokens[check_index + 1].end : tokens[close].start].strip()
source_start = tokens[constraint_start].start
source_end = tokens[close].end
return (
Check(
expression,
name=name,
column=column,
options=_parse_options(expression, column),
sql=item[source_start:source_end],
start=base_offset + source_start,
end=base_offset + source_end,
),
close + 1,
)
def _column_checks(
item: str, tokens: list[_Token], column: str, base_offset: int
) -> list[Check]:
checks: list[Check] = []
pending_name = ""
pending_start: int | None = None
index = 1
while index < len(tokens):
token = tokens[index]
if token.text == "(":
index = _matching_paren(tokens, index) + 1
continue
if token.is_keyword("CONSTRAINT"):
if index + 1 >= len(tokens):
raise ParseError("CONSTRAINT is missing its name")
pending_name = _unquote(tokens[index + 1].text)
pending_start = index
index += 2
continue
if token.is_keyword("CHECK"):
check, index = _check_after(
item,
tokens,
index,
pending_name,
column,
pending_start if pending_start is not None else index,
base_offset,
)
checks.append(check)
pending_name = ""
pending_start = None
continue
if (
token.kind == "word"
and token.text.upper() in _OTHER_COLUMN_CONSTRAINT_KEYWORDS
):
pending_name = ""
pending_start = None
index += 1
return checks
def _table_body(create_sql: str) -> tuple[str, int] | None:
all_tokens = _lex(create_sql)
tokens = _meaningful(all_tokens)
if not tokens or not tokens[0].is_keyword("CREATE"):
raise ParseError("Expected CREATE TABLE")
index = 1
if index < len(tokens) and (
tokens[index].is_keyword("TEMP") or tokens[index].is_keyword("TEMPORARY")
):
index += 1
if index < len(tokens) and tokens[index].is_keyword("VIRTUAL"):
return None
if index >= len(tokens) or not tokens[index].is_keyword("TABLE"):
raise ParseError("Expected CREATE TABLE")
index += 1
if (
index + 2 < len(tokens)
and tokens[index].is_keyword("IF")
and tokens[index + 1].is_keyword("NOT")
and tokens[index + 2].is_keyword("EXISTS")
):
index += 3
if index >= len(tokens):
raise ParseError("CREATE TABLE is missing its table name")
index += 1
if index + 1 < len(tokens) and tokens[index].text == ".":
index += 2
if index < len(tokens) and tokens[index].is_keyword("AS"):
return None
if index >= len(tokens) or tokens[index].text != "(":
raise ParseError("CREATE TABLE is missing its column list")
close = _matching_paren(tokens, index)
trailing = tokens[close + 1 :]
allowed_trailing = {"STRICT", "WITHOUT", "ROWID", ",", ";"}
if any(token.text.upper() not in allowed_trailing for token in trailing):
raise ParseError("Unexpected SQL after CREATE TABLE column list")
body_start = tokens[index].end
body_end = tokens[close].start
return create_sql[body_start:body_end], body_start
def parse_checks(create_sql: str) -> list[Check]:
"""Return CHECK constraints from a valid SQLite CREATE TABLE statement."""
body_info = _table_body(create_sql)
if body_info is None:
return []
body, body_start = body_info
body_tokens = _lex(body)
checks: list[Check] = []
for item, item_start, _ in _split_spans(body, body_tokens):
item_tokens = _meaningful(_lex(item))
if not item_tokens:
continue
item_index = 0
constraint_name = ""
if item_tokens[item_index].is_keyword("CONSTRAINT"):
if len(item_tokens) < 2:
raise ParseError("CONSTRAINT is missing its name")
constraint_name = _unquote(item_tokens[1].text)
item_index = 2
head = item_tokens[item_index] if item_index < len(item_tokens) else None
if (
head
and head.kind == "word"
and head.text.upper() in _TABLE_CONSTRAINT_KEYWORDS
):
if head.is_keyword("CHECK"):
check, _ = _check_after(
item,
item_tokens,
item_index,
constraint_name,
"",
0,
body_start + item_start,
)
checks.append(check)
continue
column = _unquote(item_tokens[0].text)
checks.extend(
_column_checks(item, item_tokens, column, body_start + item_start)
)
return checks
def parse_column_comments(create_sql: str) -> dict[str, ColumnComments]:
"""Return comments immediately before and after each column definition."""
body_info = _table_body(create_sql)
if body_info is None:
return {}
body, _ = body_info
comments: dict[str, ColumnComments] = {}
for item, _, _ in _split_spans(body, _lex(body)):
item_tokens = _meaningful(_lex(item))
if not item_tokens:
continue
item_index = 0
if item_tokens[item_index].is_keyword("CONSTRAINT"):
item_index = 2
head = item_tokens[item_index] if item_index < len(item_tokens) else None
if (
head
and head.kind == "word"
and head.text.upper() in _TABLE_CONSTRAINT_KEYWORDS
):
continue
column = _unquote(item_tokens[0].text)
before = item[: item_tokens[0].start].strip()
after = item[item_tokens[-1].end :].strip()
if before or after:
comments[column] = ColumnComments(before=before, after=after)
return comments
def _is_identifier_token(tokens: list[_Token], index: int) -> bool:
token = tokens[index]
if index + 1 < len(tokens) and tokens[index + 1].text in ("(", "."):
return False
if index and (
tokens[index - 1].is_keyword("COLLATE") or tokens[index - 1].is_keyword("AS")
):
return False
if token.kind == "identifier":
return True
if token.kind != "word" or token.text.upper() in _SQLITE_KEYWORDS:
return False
return True
def check_references_identifier(expression: str, identifier: str) -> bool:
tokens = _meaningful(_lex(expression))
folded = _ascii_fold(identifier)
return any(
_is_identifier_token(tokens, index)
and _ascii_fold(_unquote(token.text)) == folded
for index, token in enumerate(tokens)
)
def sql_ends_in_line_comment(sql: str) -> bool:
"""Return True if appended SQL would be swallowed by a ``--`` comment."""
tokens = _lex(sql)
if not tokens:
return False
final = tokens[-1]
return (
final.kind == "comment"
and final.text.startswith("--")
and not final.text.endswith(("\n", "\r"))
)
def _valid_bare_identifier(identifier: str) -> bool:
if not identifier or identifier.upper() in _SQLITE_KEYWORDS:
return False
first = identifier[0]
if not (first == "_" or first.isalpha() or ord(first) >= 0x80):
return False
return all(
char == "_" or char == "$" or char.isalnum() or ord(char) >= 0x80
for char in identifier[1:]
)
def _quote_replacement(original: str, replacement: str) -> str:
if original.startswith('"'):
return '"{}"'.format(replacement.replace('"', '""'))
if original.startswith("`"):
return "`{}`".format(replacement.replace("`", "``"))
if original.startswith("[") and "]" not in replacement:
return f"[{replacement}]"
if _valid_bare_identifier(replacement):
return replacement
return '"{}"'.format(replacement.replace('"', '""'))
def rewrite_check_expression(expression: str, rename: dict[str, str]) -> str:
"""Rewrite column identifiers in a CHECK expression, preserving trivia."""
if not rename:
return expression
tokens = _lex(expression)
meaningful = _meaningful(tokens)
replacements = {_ascii_fold(key): value for key, value in rename.items()}
edits: list[tuple[int, int, str]] = []
for index, token in enumerate(meaningful):
if not _is_identifier_token(meaningful, index):
continue
replacement = replacements.get(_ascii_fold(_unquote(token.text)))
if replacement is not None:
edits.append(
(token.start, token.end, _quote_replacement(token.text, replacement))
)
for start, end, replacement in reversed(edits):
expression = expression[:start] + replacement + expression[end:]
return expression

1323
sqlite_utils/db.py
View file

File diff suppressed because it is too large Load diff

3
sqlite_utils/hookspecs.py
View file

@ -1,8 +1,7 @@
import sqlite3
import click
from pluggy import HookimplMarker
from pluggy import HookspecMarker
from pluggy import HookimplMarker, HookspecMarker
hookspec = HookspecMarker("sqlite_utils")
hookimpl = HookimplMarker("sqlite_utils")

27
sqlite_utils/migrations.py
View file

@ -1,19 +1,28 @@
from collections.abc import Iterable
from dataclasses import dataclass
import datetime
from typing import Callable, cast, TYPE_CHECKING
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Protocol, TypeVar, cast
if TYPE_CHECKING:
from sqlite_utils.db import Database, Table
class _MigrationFunction(Protocol):
__name__: str
def __call__(self, db: "Database", /) -> None: ...
_MigrationFunctionT = TypeVar("_MigrationFunctionT", bound=_MigrationFunction)
class Migrations:
migrations_table = "_sqlite_migrations"
@dataclass
class _Migration:
name: str
fn: Callable
fn: _MigrationFunction
transactional: bool = True
@dataclass
@ -32,7 +41,7 @@ class Migrations:
def __call__(
self, *, name: str | None = None, transactional: bool = True
) -> Callable:
) -> Callable[[_MigrationFunctionT], _MigrationFunctionT]:
"""
:param name: The name to use for this migration - if not provided,
the name of the function will be used.
@ -43,13 +52,11 @@ class Migrations:
example those that execute ``VACUUM``.
"""
def inner(func: Callable) -> Callable:
migration_name = name or getattr(func, "__name__")
def inner(func: _MigrationFunctionT) -> _MigrationFunctionT:
migration_name = name or func.__name__
if any(m.name == migration_name for m in self._migrations):
raise ValueError(
"Migration '{}' is already registered in set '{}'".format(
migration_name, self.name
)
f"Migration '{migration_name}' is already registered in set '{self.name}'"
)
self._migrations.append(
self._Migration(migration_name, func, transactional)

10
sqlite_utils/plugins.py
View file

@ -1,7 +1,7 @@
from typing import Dict, List, Union
import sys
import pluggy
import sys
from . import hookspecs
pm: pluggy.PluginManager = pluggy.PluginManager("sqlite_utils")
@ -17,13 +17,13 @@ def ensure_plugins_loaded() -> None:
_plugins_loaded = True
def get_plugins() -> List[Dict[str, Union[str, List[str]]]]:
def get_plugins() -> list[dict[str, str | list[str]]]:
ensure_plugins_loaded()
plugins: List[Dict[str, Union[str, List[str]]]] = []
plugins: list[dict[str, str | list[str]]] = []
plugin_to_distinfo = dict(pm.list_plugin_distinfo())
for plugin in pm.get_plugins():
hookcallers = pm.get_hookcallers(plugin) or []
plugin_info: Dict[str, Union[str, List[str]]] = {
plugin_info: dict[str, str | list[str]] = {
"name": plugin.__name__,
"hooks": [h.name for h in hookcallers],
}

12
sqlite_utils/recipes.py
View file

@ -1,9 +1,9 @@
from __future__ import annotations
from typing import Callable, Optional
import json
from collections.abc import Callable
from dateutil import parser
import json
IGNORE: object = object()
SET_NULL: object = object()
@ -13,8 +13,8 @@ def parsedate(
value: str,
dayfirst: bool = False,
yearfirst: bool = False,
errors: Optional[object] = None,
) -> Optional[str]:
errors: object | None = None,
) -> str | None:
"""
Parse a date and convert it to ISO date format: yyyy-mm-dd
\b
@ -44,8 +44,8 @@ def parsedatetime(
value: str,
dayfirst: bool = False,
yearfirst: bool = False,
errors: Optional[object] = None,
) -> Optional[str]:
errors: object | None = None,
) -> str | None:
"""
Parse a datetime and convert it to ISO datetime format: yyyy-mm-ddTHH:MM:SS
\b

124
sqlite_utils/utils.py
View file

@ -9,20 +9,12 @@ import itertools
import json
import os
import sys
from collections.abc import Callable, Generator, Iterable, Iterator
from typing import (
TYPE_CHECKING,
Any,
BinaryIO,
Callable,
Dict,
Generator,
Iterable,
Iterator,
List,
Optional,
Set,
Tuple,
Type,
TYPE_CHECKING,
Generic,
TypeVar,
Union,
cast,
@ -33,8 +25,8 @@ import click
from . import recipes
if TYPE_CHECKING:
import sqlite3 # noqa: F401
from sqlite3 import dbapi2 # noqa: F401
import sqlite3
from sqlite3 import dbapi2
OperationalError = dbapi2.OperationalError
else:
@ -44,7 +36,7 @@ else:
OperationalError = dbapi2.OperationalError
except ImportError:
import sqlite3 # noqa: F401
from sqlite3 import dbapi2 # noqa: F401
from sqlite3 import dbapi2
OperationalError = dbapi2.OperationalError
@ -61,8 +53,8 @@ SPATIALITE_PATHS = (
ORIGINAL_CSV_FIELD_SIZE_LIMIT = csv.field_size_limit()
# Type alias for row dictionaries - values can be various SQLite-compatible types
RowValue = Union[None, int, float, str, bytes, bool, List[str]]
Row = Dict[str, RowValue]
RowValue = None | int | float | str | bytes | bool | list[str]
Row = dict[str, RowValue]
T = TypeVar("T")
@ -103,7 +95,7 @@ def maximize_csv_field_size_limit() -> None:
field_size_limit = int(field_size_limit / 10)
def find_spatialite() -> Optional[str]:
def find_spatialite() -> str | None:
"""
The ``find_spatialite()`` function searches for the `SpatiaLite <https://www.gaia-gis.it/fossil/libspatialite/index>`__
SQLite extension in some common places. It returns a string path to the location, or ``None`` if SpatiaLite was not found.
@ -132,9 +124,9 @@ def find_spatialite() -> Optional[str]:
def suggest_column_types(
records: Iterable[Dict[str, Any]],
) -> Dict[str, type]:
all_column_types: Dict[str, Set[type]] = {}
records: Iterable[dict[str, Any]],
) -> dict[str, type]:
all_column_types: dict[str, set[type]] = {}
for record in records:
for key, value in record.items():
all_column_types.setdefault(key, set()).add(type(value))
@ -142,9 +134,9 @@ def suggest_column_types(
def types_for_column_types(
all_column_types: Dict[str, Set[type]],
) -> Dict[str, type]:
column_types: Dict[str, type] = {}
all_column_types: dict[str, set[type]],
) -> dict[str, type]:
column_types: dict[str, type] = {}
for key, types in all_column_types.items():
# Ignore null values if at least one other type present:
if len(types) > 1:
@ -153,7 +145,7 @@ def types_for_column_types(
if {None.__class__} == types:
t = str
elif len(types) == 1:
t = list(types)[0]
t = next(iter(types))
# But if it's a subclass of list / tuple / dict, use str
# instead as we will be storing it as JSON in the table
for superclass in (list, tuple, dict):
@ -190,7 +182,7 @@ def column_affinity(column_type: str) -> type:
return float
def decode_base64_values(doc: Dict[str, Any]) -> Dict[str, Any]:
def decode_base64_values(doc: dict[str, Any]) -> dict[str, Any]:
# Looks for '{"$base64": true..., "encoded": ...}' values and decodes them
to_fix = [
k
@ -263,9 +255,9 @@ class RowError(Exception):
def _extra_key_strategy(
reader: Iterable[Dict[Optional[str], object]],
ignore_extras: Optional[bool] = False,
extras_key: Optional[str] = None,
reader: Iterable[dict[str | None, object]],
ignore_extras: bool | None = False,
extras_key: str | None = None,
) -> Iterable[Row]:
# Logic for handling CSV rows with more values than there are headings
for row in reader:
@ -279,9 +271,7 @@ def _extra_key_strategy(
yield cast(Row, row)
elif not extras_key:
extras = row.pop(None)
raise RowError(
"Row {} contained these extra values: {}".format(row, extras)
)
raise RowError(f"Row {row} contained these extra values: {extras}")
else:
extras_value = row.pop(None)
row_out = cast(Row, row)
@ -291,12 +281,12 @@ def _extra_key_strategy(
def rows_from_file(
fp: BinaryIO,
format: Optional[Format] = None,
dialect: Optional[Type[csv.Dialect]] = None,
encoding: Optional[str] = None,
ignore_extras: Optional[bool] = False,
extras_key: Optional[str] = None,
) -> Tuple[Iterable[Row], Format]:
format: Format | None = None,
dialect: type[csv.Dialect] | None = None,
encoding: str | None = None,
ignore_extras: bool | None = False,
extras_key: str | None = None,
) -> tuple[Iterable[Row], Format]:
"""
Load a sequence of dictionaries from a file-like object containing one of four different formats.
@ -355,7 +345,11 @@ def rows_from_file(
reader = csv.DictReader(decoded_fp, dialect=dialect)
else:
reader = csv.DictReader(decoded_fp)
rows = _extra_key_strategy(reader, ignore_extras, extras_key)
rows = _extra_key_strategy(
cast(Iterable[dict[str | None, object]], reader),
ignore_extras,
extras_key,
)
return _CloseableIterator(iter(rows), decoded_fp), Format.CSV
elif format == Format.TSV:
rows, _ = rows_from_file(
@ -363,7 +357,7 @@ def rows_from_file(
)
return (
_extra_key_strategy(
cast(Iterable[Dict[Optional[str], object]], rows),
cast(Iterable[dict[str | None, object]], rows),
ignore_extras,
extras_key,
),
@ -379,7 +373,9 @@ def rows_from_file(
raise TypeError(
"rows_from_file() requires a file-like object that supports peek(), such as io.BytesIO"
)
if first_bytes.startswith(b"[") or first_bytes.startswith(b"{"):
if not first_bytes:
return (), Format.CSV
if first_bytes.startswith((b"[", b"{")):
# TODO: Detect newline-JSON
return rows_from_file(buffered, format=Format.JSON)
else:
@ -393,7 +389,7 @@ def rows_from_file(
detected_format = Format.TSV if dialect.delimiter == "\t" else Format.CSV
return (
_extra_key_strategy(
cast(Iterable[Dict[Optional[str], object]], rows),
cast(Iterable[dict[str | None, object]], rows),
ignore_extras,
extras_key,
),
@ -425,9 +421,9 @@ class TypeTracker:
"""
def __init__(self) -> None:
self.trackers: Dict[str, "ValueTracker"] = {}
self.trackers: dict[str, ValueTracker] = {}
def wrap(self, iterator: Iterable[Dict[str, Any]]) -> Iterable[Dict[str, Any]]:
def wrap(self, iterator: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
"""
Use this to loop through an existing iterator, tracking the column types
as part of the iteration.
@ -441,7 +437,7 @@ class TypeTracker:
yield row
@property
def types(self) -> Dict[str, str]:
def types(self) -> dict[str, str]:
"""
A dictionary mapping column names to their detected types. This can be passed
to the ``db[table_name].transform(types=tracker.types)`` method.
@ -450,17 +446,15 @@ class TypeTracker:
class ValueTracker:
couldbe: Dict[str, Callable[[object], bool]]
couldbe: dict[str, Callable[[object], bool]]
def __init__(self) -> None:
self.couldbe = {key: getattr(self, "test_" + key) for key in self.get_tests()}
@classmethod
def get_tests(cls) -> List[str]:
def get_tests(cls) -> list[str]:
return [
key.split("test_")[-1]
for key in cls.__dict__.keys()
if key.startswith("test_")
key.split("test_")[-1] for key in cls.__dict__ if key.startswith("test_")
]
def test_integer(self, value: object) -> bool:
@ -492,7 +486,7 @@ class ValueTracker:
def evaluate(self, value: object) -> None:
if not value or not self.couldbe:
return
not_these: List[str] = []
not_these: list[str] = []
for name, test in self.couldbe.items():
if not test(value):
not_these.append(name)
@ -500,12 +494,12 @@ class ValueTracker:
del self.couldbe[key]
class NullProgressBar:
class NullProgressBar(Generic[T]):
def __init__(self, *args: Iterable[T]) -> None:
self.args = args
def __iter__(self) -> Iterator[T]:
yield from self.args[0] # type: ignore
yield from self.args[0]
def update(self, value: int) -> None:
pass
@ -524,14 +518,14 @@ def progressbar(*args: Iterable[T], **kwargs: Any) -> Generator[Any, None, None]
def _compile_code(
code: str, imports: Iterable[str], variable: str = "value"
) -> Callable[..., Any]:
globals_dict: Dict[str, Any] = {"r": recipes, "recipes": recipes}
globals_dict: dict[str, Any] = {"r": recipes, "recipes": recipes}
# Handle imports first so they're available for all approaches
for import_ in imports:
globals_dict[import_.split(".")[0]] = __import__(import_)
# If user defined a convert() function, return that
try:
exec(code, globals_dict)
exec(code, globals_dict) # noqa: S102
return cast(Callable[..., object], globals_dict["convert"])
except (AttributeError, SyntaxError, NameError, KeyError, TypeError):
pass
@ -542,20 +536,20 @@ def _compile_code(
fn = eval(code, globals_dict)
if callable(fn):
return cast(Callable[..., object], fn)
except Exception:
except Exception: # noqa: BLE001, S110
pass
# Try compiling their code as a function instead
body_variants = [code]
# If single line and no 'return', try adding the return
if "\n" not in code and not code.strip().startswith("return "):
body_variants.insert(0, "return {}".format(code))
body_variants.insert(0, f"return {code}")
code_o = None
for variant in body_variants:
new_code = ["def fn({}):".format(variable)]
new_code = [f"def fn({variable}):"]
for line in variant.split("\n"):
new_code.append(" {}".format(line))
new_code.append(f" {line}")
try:
code_o = compile("\n".join(new_code), "<string>", "exec")
break
@ -566,7 +560,7 @@ def _compile_code(
if code_o is None:
raise SyntaxError("Could not compile code")
exec(code_o, globals_dict)
exec(code_o, globals_dict) # noqa: S102
return cast(Callable[..., object], globals_dict["fn"])
@ -582,7 +576,7 @@ def chunks(sequence: Iterable[T], size: int) -> Iterable[Iterable[T]]:
yield itertools.chain([item], itertools.islice(iterator, size - 1))
def hash_record(record: Dict[str, Any], keys: Optional[Iterable[str]] = None) -> str:
def hash_record(record: dict[str, Any], keys: Iterable[str] | None = None) -> str:
"""
``record`` should be a Python dictionary. Returns a sha1 hash of the
keys and values in that record.
@ -603,7 +597,7 @@ def hash_record(record: Dict[str, Any], keys: Optional[Iterable[str]] = None) ->
:param record: Record to generate a hash for
:param keys: Subset of keys to use for that hash
"""
to_hash: Dict[str, Any] = record
to_hash: dict[str, Any] = record
if keys is not None:
to_hash = {key: record[key] for key in keys}
return hashlib.sha1(
@ -613,7 +607,7 @@ def hash_record(record: Dict[str, Any], keys: Optional[Iterable[str]] = None) ->
).hexdigest()
def dedupe_keys(keys: Iterable[str]) -> List[str]:
def dedupe_keys(keys: Iterable[str]) -> list[str]:
"""
Rename duplicates in a list of column names so every name is unique,
by appending ``_2``, ``_3``... to later occurrences - skipping any
@ -636,7 +630,7 @@ def dedupe_keys(keys: Iterable[str]) -> List[str]:
new_key = key
suffix = 2
while new_key in seen or new_key in taken:
new_key = "{}_{}".format(key, suffix)
new_key = f"{key}_{suffix}"
suffix += 1
key = new_key
seen.add(key)
@ -644,7 +638,7 @@ def dedupe_keys(keys: Iterable[str]) -> List[str]:
return result
def _flatten(d: Dict[str, Any]) -> Generator[Tuple[str, Any], None, None]:
def _flatten(d: dict[str, Any]) -> Generator[tuple[str, Any], None, None]:
for key, value in d.items():
if isinstance(value, dict):
for key2, value2 in _flatten(value):
@ -653,7 +647,7 @@ def _flatten(d: Dict[str, Any]) -> Generator[Tuple[str, Any], None, None]:
yield key, value
def flatten(row: Dict[str, Any]) -> Dict[str, Any]:
def flatten(row: dict[str, Any]) -> dict[str, Any]:
"""
Turn a nested dict e.g. ``{"a": {"b": 1}}`` into a flat dict: ``{"a_b": 1}``