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/transform-empty-string-to-null

Browse source 
# Conflicts:
#	tests/test_transform.py
This commit is contained in:
Simon Willison 2026-08-13 11:23:51 -07:00
commit 368964e003
71 changed files with 4526 additions and 2375 deletions
Download patch file Download diff file Expand all files Collapse all files

14
sqlite_utils/__init__.py
View file

@ -1,7 +1,13 @@
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 ANY, suggest_column_types
__all__ = ["Database", "Migrations", "suggest_column_types", "hookimpl", "hookspec"]
__all__ = [
"ANY",
"Database",
"Migrations",
"hookimpl",
"hookspec",
"suggest_column_types",
]

250
sqlite_utils/cli.py
View file

@ -1,17 +1,30 @@
import base64
import csv as csv_std
import difflib
from typing import Any
import click
from click_default_group import DefaultGroup
from datetime import datetime, timezone
import hashlib
import inspect
import io
import itertools
import json
import os
import pathlib
import pdb # noqa: T100
import sys
import textwrap
from datetime import datetime, timezone
from runpy import run_module
from typing import Any
import click
import tabulate
from click_default_group import DefaultGroup
import sqlite_utils
from sqlite_utils import recipes
from sqlite_utils.db import (
DEFAULT,
AlterError,
BadMultiValues,
DEFAULT,
DescIndex,
InvalidColumns,
NoTable,
@ -19,36 +32,28 @@ from sqlite_utils.db import (
PrimaryKeyRequired,
quote_identifier,
)
from sqlite_utils.plugins import ensure_plugins_loaded, pm, get_plugins
from sqlite_utils.plugins import ensure_plugins_loaded, get_plugins, pm
from sqlite_utils.utils import maximize_csv_field_size_limit
from sqlite_utils import recipes
import textwrap
import inspect
import io
import itertools
import json
import os
import pdb
import sys
import csv as csv_std
import tabulate
from .utils import (
Format,
OperationalError,
TypeTracker,
_compile_code,
chunks,
decode_base64_values,
dedupe_keys,
file_progress,
find_spatialite,
flatten as _flatten,
sqlite3,
decode_base64_values,
progressbar,
rows_from_file,
Format,
TypeTracker,
sqlite3,
)
from .utils import (
flatten as _flatten,
)
CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"])
CONTEXT_SETTINGS = {"help_option_names": ["-h", "--help"]}
def _register_db_for_cleanup(db):
@ -67,11 +72,11 @@ def _close_databases(ctx):
for db in ctx.meta.get("_databases_to_close", []):
try:
db.close()
except Exception:
except sqlite3.Error:
pass
VALID_COLUMN_TYPES = ("INTEGER", "TEXT", "FLOAT", "REAL", "BLOB")
VALID_COLUMN_TYPES = ("INTEGER", "TEXT", "FLOAT", "REAL", "BLOB", "ANY")
UNICODE_ERROR = """
{}
@ -174,7 +179,6 @@ def functions_option(fn):
@click.version_option()
def cli():
"Commands for interacting with a SQLite database"
pass
@cli.command()
@ -485,7 +489,17 @@ def dump(path, load_extension):
@click.argument(
"col_type",
type=click.Choice(
["integer", "int", "float", "real", "text", "str", "blob", "bytes"],
[
"integer",
"int",
"float",
"real",
"text",
"str",
"blob",
"bytes",
"any",
],
case_sensitive=False,
),
required=False,
@ -891,7 +905,7 @@ def enable_counts(path, tables, load_extension):
# Check all tables exist
bad_tables = [table for table in tables if not db[table].exists()]
if bad_tables:
raise click.ClickException("Invalid tables: {}".format(bad_tables))
raise click.ClickException(f"Invalid tables: {bad_tables}")
for table in tables:
db.table(table).enable_counts()
@ -1089,7 +1103,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,
@ -1140,9 +1154,7 @@ def insert_upsert_implementation(
)
):
raise click.ClickException(
"{}\n\nTry using --alter to add additional columns".format(
e.args[0]
)
f"{e.args[0]}\n\nTry using --alter to add additional columns"
)
# If we can find sql= and parameters= arguments, show those
variables = _find_variables(e.__traceback__, ["sql", "parameters"])
@ -1240,7 +1252,7 @@ def insert_upsert_implementation(
reader = csv_std.reader(decoded, **csv_reader_args) # type: ignore
first_row = next(reader)
if no_headers:
headers = ["untitled_{}".format(i + 1) for i in range(len(first_row))]
headers = [f"untitled_{i + 1}" for i in range(len(first_row))]
reader = itertools.chain([first_row], reader)
else:
headers = first_row
@ -1269,9 +1281,7 @@ def insert_upsert_implementation(
docs = [docs]
except json.decoder.JSONDecodeError as ex:
raise click.ClickException(
"Invalid JSON - use --csv for CSV or --tsv for TSV files\n\nJSON error: {}".format(
ex
)
f"Invalid JSON - use --csv for CSV or --tsv for TSV files\n\nJSON error: {ex}"
)
if flatten:
docs = (_flatten(doc) for doc in docs)
@ -1290,7 +1300,7 @@ def insert_upsert_implementation(
docs = (fn(doc["line"]) for doc in docs)
elif text:
# Special case: this is allowed to be an iterable
text_value = list(docs)[0]["text"]
text_value = next(iter(docs))["text"]
fn_return = fn(text_value)
if isinstance(fn_return, dict):
docs = [fn_return]
@ -1758,7 +1768,7 @@ def create_table(
height real \\
photo blob --pk id
Valid column types are text, integer, real, float and blob.
Valid column types are text, integer, real, float, blob and any.
"""
db = sqlite_utils.Database(path)
_register_db_for_cleanup(db)
@ -1774,17 +1784,14 @@ def create_table(
ctype = columns.pop(0)
if ctype.upper() not in VALID_COLUMN_TYPES:
raise click.ClickException(
"column types must be one of {}".format(VALID_COLUMN_TYPES)
f"column types must be one of {VALID_COLUMN_TYPES}"
)
coltypes[name] = ctype.upper()
# Does table already exist?
if table in db.table_names():
if not ignore and not replace and not transform:
raise click.ClickException(
'Table "{}" already exists. Use --replace to delete and replace it.'.format(
table
)
)
if table in db.table_names() and not ignore and not replace and not transform:
raise click.ClickException(
f'Table "{table}" already exists. Use --replace to delete and replace it.'
)
db.table(table).create(
coltypes,
pk=pks[0] if len(pks) == 1 else pks,
@ -1819,7 +1826,7 @@ def duplicate(path, table, new_table, ignore, load_extension):
db.table(table).duplicate(new_table)
except NoTable:
if not ignore:
raise click.ClickException('Table "{}" does not exist'.format(table))
raise click.ClickException(f'Table "{table}" does not exist')
@cli.command(name="rename-table")
@ -1843,9 +1850,7 @@ def rename_table(path, table, new_name, ignore, load_extension):
db.rename_table(table, new_name)
except sqlite3.OperationalError as ex:
if not ignore:
raise click.ClickException(
'Table "{}" could not be renamed. {}'.format(table, str(ex))
)
raise click.ClickException(f'Table "{table}" could not be renamed. {ex!s}')
@cli.command(name="drop-table")
@ -1874,10 +1879,10 @@ def drop_table(path, table, ignore, load_extension):
# A view exists with this name
if not ignore:
raise click.ClickException(
'"{}" is a view, not a table - use drop-view to drop it'.format(table)
f'"{table}" is a view, not a table - use drop-view to drop it'
)
except OperationalError:
raise click.ClickException('Table "{}" does not exist'.format(table))
raise click.ClickException(f'Table "{table}" does not exist')
@cli.command(name="create-view")
@ -1919,9 +1924,7 @@ def create_view(path, view, select, ignore, replace, load_extension):
db.view(view).drop()
else:
raise click.ClickException(
'View "{}" already exists. Use --replace to delete and replace it.'.format(
view
)
f'View "{view}" already exists. Use --replace to delete and replace it.'
)
db.create_view(view, select)
@ -1953,9 +1956,9 @@ def drop_view(path, view, ignore, load_extension):
return
if view in db.table_names():
raise click.ClickException(
'"{}" is a table, not a view - use drop-table to drop it'.format(view)
f'"{view}" is a table, not a view - use drop-table to drop it'
)
raise click.ClickException('View "{}" does not exist'.format(view))
raise click.ClickException(f'View "{view}" does not exist')
@cli.command()
@ -2177,7 +2180,7 @@ def memory(
file_path = pathlib.Path(path)
stem = file_path.stem
if stem_counts.get(stem):
file_table = "{}_{}".format(stem, stem_counts[stem])
file_table = f"{stem}_{stem_counts[stem]}"
else:
file_table = stem
stem_counts[stem] = stem_counts.get(stem, 1) + 1
@ -2196,14 +2199,14 @@ def memory(
if tracker is not None and db.table(file_table).exists():
db.table(file_table).transform(types=tracker.types)
# Add convenient t / t1 / t2 views
view_names = ["t{}".format(i + 1)]
view_names = [f"t{i + 1}"]
if i == 0:
view_names.append("t")
for view_name in view_names:
if not db[view_name].exists():
db.create_view(
view_name,
"select * from {}".format(quote_identifier(file_table)),
f"select * from {quote_identifier(file_table)}",
)
finally:
if should_close_fp and fp:
@ -2373,19 +2376,17 @@ def search(
# Check table exists
table_obj = db.table(dbtable)
if not table_obj.exists():
raise click.ClickException("Table '{}' does not exist".format(dbtable))
raise click.ClickException(f"Table '{dbtable}' does not exist")
if not table_obj.detect_fts():
raise click.ClickException(
"Table '{}' is not configured for full-text search".format(dbtable)
f"Table '{dbtable}' is not configured for full-text search"
)
if column:
# Check they all exist
table_columns = table_obj.columns_dict
for c in column:
if c not in table_columns:
raise click.ClickException(
"Table '{}' has no column '{}".format(dbtable, c)
)
raise click.ClickException(f"Table '{dbtable}' has no column '{c}")
sql = table_obj.search_sql(columns=column, order_by=order, limit=limit)
if show_sql:
click.echo(sql)
@ -2412,7 +2413,7 @@ def search(
except click.ClickException as e:
if "malformed MATCH expression" in str(e) or "unterminated string" in str(e):
raise click.ClickException(
"{}\n\nTry running this again with the --quote option".format(str(e))
f"{e!s}\n\nTry running this again with the --quote option"
)
else:
raise
@ -2479,15 +2480,17 @@ def rows(
columns = "*"
if column:
columns = ", ".join(quote_identifier(c) for c in column)
sql = "select {} from {}".format(columns, quote_identifier(dbtable))
sql = f"select {columns} from {quote_identifier(dbtable)}"
if where:
sql += " where " + where
if order:
sql += " order by " + order
if limit:
sql += " limit {}".format(limit)
sql += f" limit {limit}"
if offset:
sql += " offset {}".format(offset)
if not limit:
sql += " limit -1"
sql += f" offset {offset}"
ctx.invoke(
query,
path=path,
@ -2675,12 +2678,10 @@ def schema(
"--type",
type=(
str,
click.Choice(
["INTEGER", "TEXT", "FLOAT", "REAL", "BLOB"], case_sensitive=False
),
click.Choice(list(VALID_COLUMN_TYPES), case_sensitive=False),
),
multiple=True,
help="Change column type to INTEGER, TEXT, FLOAT, REAL or BLOB",
help="Change column type to INTEGER, TEXT, FLOAT, REAL, BLOB or ANY",
)
@click.option("--drop", type=str, multiple=True, help="Drop this column")
@click.option(
@ -2760,7 +2761,7 @@ def transform(
for column, ctype in type:
if ctype.upper() not in VALID_COLUMN_TYPES:
raise click.ClickException(
"column types must be one of {}".format(VALID_COLUMN_TYPES)
f"column types must be one of {VALID_COLUMN_TYPES}"
)
types[column] = ctype.upper()
@ -2858,12 +2859,12 @@ def extract(
db = sqlite_utils.Database(path)
_register_db_for_cleanup(db)
_load_extensions(db, load_extension)
kwargs: dict[str, Any] = dict(
columns=columns,
table=other_table,
fk_column=fk_column,
rename=dict(rename),
)
kwargs: dict[str, Any] = {
"columns": columns,
"table": other_table,
"fk_column": fk_column,
"rename": dict(rename),
}
try:
db.table(table).extract(**kwargs)
except (NoTable, InvalidColumns) as e:
@ -2958,7 +2959,7 @@ def insert_files(
with progressbar(paths_and_relative_paths, silent=silent) as bar:
def to_insert():
for path, relative_path in bar:
for file_path, relative_path in bar:
row = {}
# content_text is special case as it considers 'encoding'
@ -2970,19 +2971,21 @@ def insert_files(
raise UnicodeDecodeErrorForPath(e, resolved)
lookups = dict(FILE_COLUMNS, content_text=_content_text)
if path == "-":
if file_path == "-":
stdin_data = sys.stdin.buffer.read()
# We only support a subset of columns for this case
lookups = {
"name": lambda p: name or "-",
"path": lambda p: name or "-",
"content": lambda p: stdin_data,
"content_text": lambda p: stdin_data.decode(
"content": lambda p, data=stdin_data: data,
"content_text": lambda p, data=stdin_data: data.decode(
encoding or "utf-8"
),
"sha256": lambda p: hashlib.sha256(stdin_data).hexdigest(),
"md5": lambda p: hashlib.md5(stdin_data).hexdigest(),
"size": lambda p: len(stdin_data),
"sha256": lambda p, data=stdin_data: hashlib.sha256(
data
).hexdigest(),
"md5": lambda p, data=stdin_data: hashlib.md5(data).hexdigest(),
"size": lambda p, data=stdin_data: len(data),
}
for coldef in column:
if ":" in coldef:
@ -2990,7 +2993,7 @@ def insert_files(
else:
colname, coltype = coldef, coldef
try:
value = lookups[coltype](path)
value = lookups[coltype](file_path)
row[colname] = value
except KeyError:
raise click.ClickException(
@ -3018,7 +3021,7 @@ def insert_files(
except UnicodeDecodeErrorForPath as e:
raise click.ClickException(
UNICODE_ERROR.format(
"Could not read file '{}' as text\n\n{}".format(e.path, e.exception)
f"Could not read file '{e.path}' as text\n\n{e.exception}"
)
)
@ -3196,7 +3199,7 @@ def _generate_convert_help():
for name in recipe_names:
fn = getattr(recipes, name)
doc = textwrap.dedent(fn.__doc__.rstrip()).replace("\b\n", "")
help += "\n\nr.{}{}\n\n\b{}".format(name, str(inspect.signature(fn)), doc)
help += f"\n\nr.{name}{inspect.signature(fn)!s}\n\n\b{doc}"
help += "\n\n"
help += textwrap.dedent("""
You can use these recipes like so:
@ -3280,26 +3283,21 @@ 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)
sql = """
select
[{column}] as value,
preview_transform([{column}]) as preview
from [{table}]{where} limit 10
{column} as value,
preview_transform({column}) as preview
from {table}{where} limit 10
""".format(
column=columns[0],
table=table,
where=" where {}".format(where) if where is not None else "",
column=quote_identifier(columns[0]),
table=quote_identifier(table),
where=f" where {where}" if where is not None else "",
)
for row in db.conn.execute(sql, where_args).fetchall():
click.echo(str(row[0]))
@ -3319,7 +3317,7 @@ def convert(
def wrapped_fn(value):
try:
return fn_(value)
except Exception as ex:
except Exception as ex: # noqa: BLE001
print("\nException raised, dropping into pdb...:", ex)
pdb.post_mortem(ex.__traceback__)
sys.exit(1)
@ -3339,9 +3337,7 @@ def convert(
)
except BadMultiValues as e:
raise click.ClickException(
"When using --multi code must return a Python dictionary - returned: {}".format(
repr(e.values)
)
f"When using --multi code must return a Python dictionary - returned: {e.values!r}"
)
@ -3459,7 +3455,7 @@ def create_spatial_index(db_path, table, column_name, load_extension):
def _find_migration_files(migrations):
if not migrations:
migrations = [pathlib.Path(".").resolve()]
migrations = [pathlib.Path.cwd()]
files = set()
for path_str in migrations:
path = pathlib.Path(path_str)
@ -3484,7 +3480,7 @@ def _load_migration_sets(files):
"__file__": str(filepath),
"__name__": "__sqlite_utils_migration__",
}
exec(code, namespace)
exec(code, namespace) # noqa: S102
migration_sets.extend(
obj for obj in namespace.values() if _compatible_migration_set(obj)
)
@ -3493,17 +3489,17 @@ def _load_migration_sets(files):
def _display_migration_list(db, migration_sets):
for migration_set in migration_sets:
click.echo("Migrations for: {}".format(migration_set.name))
click.echo(f"Migrations for: {migration_set.name}")
click.echo()
click.echo(" Applied:")
for migration in migration_set.applied(db):
click.echo(" {} - {}".format(migration.name, migration.applied_at))
click.echo(f" {migration.name} - {migration.applied_at}")
click.echo()
click.echo(" Pending:")
output = False
for migration in migration_set.pending(db):
output = True
click.echo(" {}".format(migration.name))
click.echo(f" {migration.name}")
if not output:
click.echo(" (none)")
click.echo()
@ -3583,7 +3579,7 @@ def migrate(db_path, migrations, stop_before, list_, verbose):
prev_schema = db.schema
if verbose:
click.echo("Migrating {}".format(db_path))
click.echo(f"Migrating {db_path}")
click.echo("\nSchema before:\n")
click.echo(textwrap.indent(prev_schema, " ") or " (empty)")
click.echo()
@ -3594,9 +3590,7 @@ def migrate(db_path, migrations, stop_before, list_, verbose):
names = {m.name for m in migration_set.pending(db)}
names.update(m.name for m in migration_set.applied(db))
known_names.update(names)
known_names.update(
"{}:{}".format(migration_set.name, name) for name in names
)
known_names.update(f"{migration_set.name}:{name}" for name in names)
unknown = [value for value in stop_before if value not in known_names]
if unknown:
raise click.ClickException(
@ -3652,7 +3646,7 @@ def _render_common(title, values):
return ""
lines = [title]
for value, count in values:
lines.append(" {}: {}".format(count, value))
lines.append(f" {count}: {value}")
return "\n".join(lines)
@ -3722,7 +3716,7 @@ def maybe_json(value):
if not isinstance(value, str):
return value
stripped = value.strip()
if not (stripped.startswith("{") or stripped.startswith("[")):
if not (stripped.startswith(("{", "["))):
return value
try:
return json.loads(stripped)
@ -3740,7 +3734,7 @@ def json_binary(value):
def verify_is_dict(doc):
if not isinstance(doc, dict):
raise click.ClickException(
"Rows must all be dictionaries, got: {}".format(repr(doc)[:1000])
f"Rows must all be dictionaries, got: {repr(doc)[:1000]}"
)
return doc
@ -3768,14 +3762,14 @@ def _register_functions(db, functions):
try:
functions = pathlib.Path(functions).read_text()
except FileNotFoundError:
raise click.ClickException("File not found: {}".format(functions))
raise click.ClickException(f"File not found: {functions}")
sqlite3.enable_callback_tracebacks(True)
globals = {}
try:
exec(functions, globals)
exec(functions, globals) # noqa: S102
except SyntaxError as ex:
raise click.ClickException("Error in functions definition: {}".format(ex))
raise click.ClickException(f"Error in functions definition: {ex}")
# Register all callables in the locals dict:
for name, value in globals.items():
if callable(value) and not name.startswith("_"):
@ -3796,13 +3790,13 @@ def _rows_from_code(code):
try:
code = pathlib.Path(code).read_text()
except FileNotFoundError:
raise click.ClickException("File not found: {}".format(code))
namespace = {}
raise click.ClickException(f"File not found: {code}")
namespace: dict[str, Any] = {}
try:
exec(code, namespace)
exec(code, namespace) # noqa: S102
except SyntaxError as ex:
raise click.ClickException("Error in --code: {}".format(ex))
rows = namespace.get("rows")
raise click.ClickException(f"Error in --code: {ex}")
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

1400
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

130
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,12 +53,16 @@ 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")
class ANY:
"""Marker type for an SQLite ``ANY`` column."""
class _CloseableIterator(Iterator[Row]):
"""Iterator wrapper that closes a file when iteration is complete."""
@ -103,7 +99,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 +128,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 +138,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 +149,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):
@ -186,11 +182,13 @@ def column_affinity(column_type: str) -> type:
return bytes
if "REAL" in column_type or "FLOA" in column_type or "DOUB" in column_type:
return float
if column_type == "ANY":
return ANY
# Default is 'NUMERIC', which we currently also treat as float
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 +261,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 +277,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 +287,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 +351,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 +363,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 +379,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 +395,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 +427,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 +443,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 +452,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 +492,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 +500,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 +524,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 +542,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 +566,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 +582,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 +603,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 +613,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 +636,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 +644,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 +653,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}``