mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-07-24 01:44:31 +02:00
Compare commits
2 commits
main
...
claude/add
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9b83ae0325 |
||
|
|
83c50527aa |
7 changed files with 234 additions and 129 deletions
30
mypy.ini
30
mypy.ini
|
|
@ -1,4 +1,32 @@
|
|||
[mypy]
|
||||
python_version = 3.10
|
||||
warn_return_any = False
|
||||
warn_unused_configs = True
|
||||
warn_redundant_casts = False
|
||||
warn_unused_ignores = False
|
||||
check_untyped_defs = True
|
||||
disallow_untyped_defs = False
|
||||
disallow_incomplete_defs = False
|
||||
no_implicit_optional = True
|
||||
strict_equality = True
|
||||
|
||||
[mypy-pysqlite3,sqlean,sqlite_dump]
|
||||
[mypy-sqlite_utils.cli]
|
||||
ignore_errors = True
|
||||
|
||||
[mypy-pysqlite3.*]
|
||||
ignore_missing_imports = True
|
||||
|
||||
[mypy-sqlean.*]
|
||||
ignore_missing_imports = True
|
||||
|
||||
[mypy-sqlite_dump.*]
|
||||
ignore_missing_imports = True
|
||||
|
||||
[mypy-sqlite_fts4.*]
|
||||
ignore_missing_imports = True
|
||||
|
||||
[mypy-pandas.*]
|
||||
ignore_missing_imports = True
|
||||
|
||||
[mypy-numpy.*]
|
||||
ignore_missing_imports = True
|
||||
|
|
@ -22,7 +22,7 @@ import os
|
|||
import pathlib
|
||||
import re
|
||||
import secrets
|
||||
from sqlite_fts4 import rank_bm25 # type: ignore
|
||||
from sqlite_fts4 import rank_bm25
|
||||
import textwrap
|
||||
from typing import (
|
||||
cast,
|
||||
|
|
@ -32,10 +32,13 @@ from typing import (
|
|||
Generator,
|
||||
Iterable,
|
||||
Sequence,
|
||||
Set,
|
||||
Type,
|
||||
Union,
|
||||
Optional,
|
||||
List,
|
||||
Tuple,
|
||||
TYPE_CHECKING,
|
||||
)
|
||||
import uuid
|
||||
from sqlite_utils.plugins import pm
|
||||
|
|
@ -81,14 +84,14 @@ def quote_identifier(identifier: str) -> str:
|
|||
|
||||
|
||||
try:
|
||||
import pandas as pd # type: ignore
|
||||
import pandas as pd
|
||||
except ImportError:
|
||||
pd = None # type: ignore
|
||||
pd = None
|
||||
|
||||
try:
|
||||
import numpy as np # type: ignore
|
||||
import numpy as np
|
||||
except ImportError:
|
||||
np = None # type: ignore
|
||||
np = None
|
||||
|
||||
Column = namedtuple(
|
||||
"Column", ("cid", "name", "type", "notnull", "default_value", "is_pk")
|
||||
|
|
@ -245,7 +248,7 @@ if np:
|
|||
|
||||
# If pandas is available, add more types
|
||||
if pd:
|
||||
COLUMN_TYPE_MAPPING.update({pd.Timestamp: "TEXT"}) # type: ignore
|
||||
COLUMN_TYPE_MAPPING.update({pd.Timestamp: "TEXT"})
|
||||
|
||||
|
||||
class AlterError(Exception):
|
||||
|
|
@ -287,7 +290,7 @@ class DescIndex(str):
|
|||
class BadMultiValues(Exception):
|
||||
"With multi=True code must return a Python dictionary"
|
||||
|
||||
def __init__(self, values):
|
||||
def __init__(self, values: Any) -> None:
|
||||
self.values = values
|
||||
|
||||
|
||||
|
|
@ -385,7 +388,12 @@ class Database:
|
|||
def __enter__(self) -> "Database":
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
||||
def __exit__(
|
||||
self,
|
||||
exc_type: Optional[Type[BaseException]],
|
||||
exc_val: Optional[BaseException],
|
||||
exc_tb: Optional[Any],
|
||||
) -> None:
|
||||
self.close()
|
||||
|
||||
def close(self) -> None:
|
||||
|
|
@ -393,7 +401,7 @@ class Database:
|
|||
self.conn.close()
|
||||
|
||||
@contextlib.contextmanager
|
||||
def ensure_autocommit_off(self):
|
||||
def ensure_autocommit_off(self) -> Generator[None, None, None]:
|
||||
"""
|
||||
Ensure autocommit is off for this database connection.
|
||||
|
||||
|
|
@ -412,7 +420,9 @@ class Database:
|
|||
self.conn.isolation_level = old_isolation_level
|
||||
|
||||
@contextlib.contextmanager
|
||||
def tracer(self, tracer: Optional[Callable] = None):
|
||||
def tracer(
|
||||
self, tracer: Optional[Callable[..., Any]] = None
|
||||
) -> Generator["Database", None, None]:
|
||||
"""
|
||||
Context manager to temporarily set a tracer function - all executed SQL queries will
|
||||
be passed to this.
|
||||
|
|
@ -451,11 +461,11 @@ class Database:
|
|||
|
||||
def register_function(
|
||||
self,
|
||||
fn: Optional[Callable] = None,
|
||||
fn: Optional[Callable[..., Any]] = None,
|
||||
deterministic: bool = False,
|
||||
replace: bool = False,
|
||||
name: Optional[str] = None,
|
||||
):
|
||||
) -> Optional[Callable[[Callable[..., Any]], Callable[..., Any]]]:
|
||||
"""
|
||||
``fn`` will be made available as a function within SQL, with the same name and number
|
||||
of arguments. Can be used as a decorator::
|
||||
|
|
@ -478,12 +488,12 @@ class Database:
|
|||
:param name: name of the SQLite function - if not specified, the Python function name will be used
|
||||
"""
|
||||
|
||||
def register(fn):
|
||||
def register(fn: Callable[..., Any]) -> Callable[..., Any]:
|
||||
fn_name = name or fn.__name__
|
||||
arity = len(inspect.signature(fn).parameters)
|
||||
if not replace and (fn_name, arity) in self._registered_functions:
|
||||
return fn
|
||||
kwargs = {}
|
||||
kwargs: Dict[str, Any] = {}
|
||||
registered = False
|
||||
if deterministic:
|
||||
# Try this, but fall back if sqlite3.NotSupportedError
|
||||
|
|
@ -503,12 +513,13 @@ class Database:
|
|||
return register
|
||||
else:
|
||||
register(fn)
|
||||
return None
|
||||
|
||||
def register_fts4_bm25(self):
|
||||
def register_fts4_bm25(self) -> None:
|
||||
"Register the ``rank_bm25(match_info)`` function used for calculating relevance with SQLite FTS4."
|
||||
self.register_function(rank_bm25, deterministic=True, replace=True)
|
||||
|
||||
def attach(self, alias: str, filepath: Union[str, pathlib.Path]):
|
||||
def attach(self, alias: str, filepath: Union[str, pathlib.Path]) -> None:
|
||||
"""
|
||||
Attach another SQLite database file to this connection with the specified alias, equivalent to::
|
||||
|
||||
|
|
@ -525,8 +536,8 @@ class Database:
|
|||
self.execute(attach_sql)
|
||||
|
||||
def query(
|
||||
self, sql: str, params: Optional[Union[Iterable, dict]] = None
|
||||
) -> Generator[dict, None, None]:
|
||||
self, sql: str, params: Optional[Union[Iterable[Any], Dict[str, Any]]] = None
|
||||
) -> Generator[Dict[str, Any], None, None]:
|
||||
"""
|
||||
Execute ``sql`` and return an iterable of dictionaries representing each row.
|
||||
|
||||
|
|
@ -566,7 +577,7 @@ class Database:
|
|||
self._tracer(sql, None)
|
||||
return self.conn.executescript(sql)
|
||||
|
||||
def table(self, table_name: str, **kwargs) -> "Table":
|
||||
def table(self, table_name: str, **kwargs: Any) -> "Table":
|
||||
"""
|
||||
Return a table object, optionally configured with default options.
|
||||
|
||||
|
|
@ -598,11 +609,12 @@ class Database:
|
|||
# Normally we would use .execute(sql, [params]) for escaping, but
|
||||
# occasionally that isn't available - most notable when we need
|
||||
# to include a "... DEFAULT 'value'" in a column definition.
|
||||
return self.execute(
|
||||
result: str = self.execute(
|
||||
# Use SQLite itself to correctly escape this string:
|
||||
"SELECT quote(:value)",
|
||||
{"value": value},
|
||||
).fetchone()[0]
|
||||
return result
|
||||
|
||||
def quote_fts(self, query: str) -> str:
|
||||
"""
|
||||
|
|
@ -763,9 +775,10 @@ class Database:
|
|||
|
||||
https://www.sqlite.org/pragma.html#pragma_journal_mode
|
||||
"""
|
||||
return self.execute("PRAGMA journal_mode;").fetchone()[0]
|
||||
result: str = self.execute("PRAGMA journal_mode;").fetchone()[0]
|
||||
return result
|
||||
|
||||
def enable_wal(self):
|
||||
def enable_wal(self) -> None:
|
||||
"""
|
||||
Sets ``journal_mode`` to ``'wal'`` to enable Write-Ahead Log mode.
|
||||
"""
|
||||
|
|
@ -773,17 +786,17 @@ class Database:
|
|||
with self.ensure_autocommit_off():
|
||||
self.execute("PRAGMA journal_mode=wal;")
|
||||
|
||||
def disable_wal(self):
|
||||
def disable_wal(self) -> None:
|
||||
"Sets ``journal_mode`` back to ``'delete'`` to disable Write-Ahead Log mode."
|
||||
if self.journal_mode != "delete":
|
||||
with self.ensure_autocommit_off():
|
||||
self.execute("PRAGMA journal_mode=delete;")
|
||||
|
||||
def _ensure_counts_table(self):
|
||||
def _ensure_counts_table(self) -> None:
|
||||
with self.conn:
|
||||
self.execute(_COUNTS_TABLE_CREATE_SQL.format(self._counts_table_name))
|
||||
|
||||
def enable_counts(self):
|
||||
def enable_counts(self) -> None:
|
||||
"""
|
||||
Enable trigger-based count caching for every table in the database, see
|
||||
:ref:`python_api_cached_table_counts`.
|
||||
|
|
@ -812,12 +825,12 @@ class Database:
|
|||
except OperationalError:
|
||||
return {}
|
||||
|
||||
def reset_counts(self):
|
||||
def reset_counts(self) -> None:
|
||||
"Re-calculate cached counts for tables."
|
||||
tables = [table for table in self.tables if table.has_counts_triggers]
|
||||
with self.conn:
|
||||
self._ensure_counts_table()
|
||||
counts_table = self[self._counts_table_name]
|
||||
counts_table = cast("Table", self[self._counts_table_name])
|
||||
counts_table.delete_where()
|
||||
counts_table.insert_all(
|
||||
{"table": table.name, "count": table.execute_count()}
|
||||
|
|
@ -825,8 +838,8 @@ class Database:
|
|||
)
|
||||
|
||||
def execute_returning_dicts(
|
||||
self, sql: str, params: Optional[Union[Iterable, dict]] = None
|
||||
) -> List[dict]:
|
||||
self, sql: str, params: Optional[Union[Iterable[Any], Dict[str, Any]]] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
return list(self.query(sql, params))
|
||||
|
||||
def resolve_foreign_keys(
|
||||
|
|
@ -954,7 +967,7 @@ class Database:
|
|||
column_items = list(columns.items())
|
||||
if column_order is not None:
|
||||
|
||||
def sort_key(p):
|
||||
def sort_key(p: Tuple[str, Any]) -> int:
|
||||
return column_order.index(p[0]) if p[0] in column_order else 999
|
||||
|
||||
column_items.sort(key=sort_key)
|
||||
|
|
@ -1157,7 +1170,7 @@ class Database:
|
|||
hash_id_columns=hash_id_columns,
|
||||
)
|
||||
|
||||
def rename_table(self, name: str, new_name: str):
|
||||
def rename_table(self, name: str, new_name: str) -> None:
|
||||
"""
|
||||
Rename a table.
|
||||
|
||||
|
|
@ -1172,7 +1185,7 @@ class Database:
|
|||
|
||||
def create_view(
|
||||
self, name: str, sql: str, ignore: bool = False, replace: bool = False
|
||||
):
|
||||
) -> "Database":
|
||||
"""
|
||||
Create a new SQL view with the specified name - ``sql`` should start with ``SELECT ...``.
|
||||
|
||||
|
|
@ -1218,7 +1231,9 @@ class Database:
|
|||
candidates.append(table_obj.name)
|
||||
return candidates
|
||||
|
||||
def add_foreign_keys(self, foreign_keys: Iterable[Tuple[str, str, str, str]]):
|
||||
def add_foreign_keys(
|
||||
self, foreign_keys: Iterable[Tuple[str, str, str, str]]
|
||||
) -> None:
|
||||
"""
|
||||
See :ref:`python_api_add_foreign_keys`.
|
||||
|
||||
|
|
@ -1270,10 +1285,10 @@ class Database:
|
|||
|
||||
self.vacuum()
|
||||
|
||||
def index_foreign_keys(self):
|
||||
def index_foreign_keys(self) -> None:
|
||||
"Create indexes for every foreign key column on every table in the database."
|
||||
for table_name in self.table_names():
|
||||
table = self[table_name]
|
||||
table = self.table(table_name)
|
||||
existing_indexes = {
|
||||
i.columns[0] for i in table.indexes if len(i.columns) == 1
|
||||
}
|
||||
|
|
@ -1281,11 +1296,11 @@ class Database:
|
|||
if fk.column not in existing_indexes:
|
||||
table.create_index([fk.column], find_unique_name=True)
|
||||
|
||||
def vacuum(self):
|
||||
def vacuum(self) -> None:
|
||||
"Run a SQLite ``VACUUM`` against the database."
|
||||
self.execute("VACUUM;")
|
||||
|
||||
def analyze(self, name=None):
|
||||
def analyze(self, name: Optional[str] = None) -> None:
|
||||
"""
|
||||
Run ``ANALYZE`` against the entire database or a named table or index.
|
||||
|
||||
|
|
@ -1351,18 +1366,21 @@ class Database:
|
|||
|
||||
|
||||
class Queryable:
|
||||
db: "Database"
|
||||
name: str
|
||||
|
||||
def exists(self) -> bool:
|
||||
"Does this table or view exist yet?"
|
||||
return False
|
||||
|
||||
def __init__(self, db, name):
|
||||
def __init__(self, db: "Database", name: str) -> None:
|
||||
self.db = db
|
||||
self.name = name
|
||||
|
||||
def count_where(
|
||||
self,
|
||||
where: Optional[str] = None,
|
||||
where_args: Optional[Union[Iterable, dict]] = None,
|
||||
where_args: Optional[Union[Iterable[Any], Dict[str, Any]]] = None,
|
||||
) -> int:
|
||||
"""
|
||||
Executes ``SELECT count(*) FROM table WHERE ...`` and returns a count.
|
||||
|
|
@ -1374,9 +1392,10 @@ class Queryable:
|
|||
sql = "select count(*) from {}".format(quote_identifier(self.name))
|
||||
if where is not None:
|
||||
sql += " where " + where
|
||||
return self.db.execute(sql, where_args or []).fetchone()[0]
|
||||
result: int = self.db.execute(sql, where_args or []).fetchone()[0]
|
||||
return result
|
||||
|
||||
def execute_count(self):
|
||||
def execute_count(self) -> int:
|
||||
# Backwards compatibility, see https://github.com/simonw/sqlite-utils/issues/305#issuecomment-890713185
|
||||
return self.count_where()
|
||||
|
||||
|
|
@ -1386,19 +1405,19 @@ class Queryable:
|
|||
return self.count_where()
|
||||
|
||||
@property
|
||||
def rows(self) -> Generator[dict, None, None]:
|
||||
def rows(self) -> Generator[Dict[str, Any], None, None]:
|
||||
"Iterate over every dictionaries for each row in this table or view."
|
||||
return self.rows_where()
|
||||
|
||||
def rows_where(
|
||||
self,
|
||||
where: Optional[str] = None,
|
||||
where_args: Optional[Union[Iterable, dict]] = None,
|
||||
where_args: Optional[Union[Iterable[Any], Dict[str, Any]]] = None,
|
||||
order_by: Optional[str] = None,
|
||||
select: str = "*",
|
||||
limit: Optional[int] = None,
|
||||
offset: Optional[int] = None,
|
||||
) -> Generator[dict, None, None]:
|
||||
) -> Generator[Dict[str, Any], None, None]:
|
||||
"""
|
||||
Iterate over every row in this table or view that matches the specified where clause.
|
||||
|
||||
|
|
@ -1800,18 +1819,18 @@ class Table(Queryable):
|
|||
self.name,
|
||||
columns,
|
||||
pk=pk,
|
||||
foreign_keys=foreign_keys,
|
||||
column_order=column_order,
|
||||
not_null=not_null,
|
||||
defaults=defaults,
|
||||
hash_id=hash_id,
|
||||
hash_id_columns=hash_id_columns,
|
||||
extracts=extracts,
|
||||
foreign_keys=foreign_keys, # type: ignore[arg-type]
|
||||
column_order=column_order, # type: ignore[arg-type]
|
||||
not_null=not_null, # type: ignore[arg-type]
|
||||
defaults=defaults, # type: ignore[arg-type]
|
||||
hash_id=hash_id, # type: ignore[arg-type]
|
||||
hash_id_columns=hash_id_columns, # type: ignore[arg-type]
|
||||
extracts=extracts, # type: ignore[arg-type]
|
||||
if_not_exists=if_not_exists,
|
||||
replace=replace,
|
||||
ignore=ignore,
|
||||
transform=transform,
|
||||
strict=strict,
|
||||
strict=bool(strict) if strict is not DEFAULT else False,
|
||||
)
|
||||
return self
|
||||
|
||||
|
|
@ -1829,7 +1848,7 @@ class Table(Queryable):
|
|||
quote_identifier(self.name),
|
||||
)
|
||||
self.db.execute(sql)
|
||||
return self.db[new_name]
|
||||
return self.db.table(new_name)
|
||||
|
||||
def transform(
|
||||
self,
|
||||
|
|
@ -2134,7 +2153,7 @@ class Table(Queryable):
|
|||
)
|
||||
)
|
||||
table = table or "_".join(columns)
|
||||
lookup_table = self.db[table]
|
||||
lookup_table = self.db.table(table)
|
||||
fk_column = fk_column or "{}_id".format(table)
|
||||
magic_lookup_column = "{}_{}".format(fk_column, os.urandom(6).hex())
|
||||
|
||||
|
|
@ -2748,7 +2767,7 @@ class Table(Queryable):
|
|||
self.name
|
||||
)
|
||||
fts_table_quoted = quote_identifier(fts_table)
|
||||
virtual_table_using = self.db[fts_table].virtual_table_using
|
||||
virtual_table_using = self.db.table(fts_table).virtual_table_using
|
||||
sql = textwrap.dedent(
|
||||
"""
|
||||
with {original} as (
|
||||
|
|
@ -2845,7 +2864,7 @@ class Table(Queryable):
|
|||
for row in cursor:
|
||||
yield dict(zip(columns, row))
|
||||
|
||||
def value_or_default(self, key, value):
|
||||
def value_or_default(self, key: str, value: Any) -> Any:
|
||||
return self._defaults[key] if value is DEFAULT else value
|
||||
|
||||
def delete(self, pk_values: Union[list, tuple, str, int, float]) -> "Table":
|
||||
|
|
@ -3033,8 +3052,8 @@ class Table(Queryable):
|
|||
self, column, fn, drop, show_progress, where=None, where_args=None
|
||||
):
|
||||
# First we execute the function
|
||||
pk_to_values = {}
|
||||
new_column_types = {}
|
||||
pk_to_values: Dict[Any, Any] = {}
|
||||
new_column_types: Dict[str, Set[type]] = {}
|
||||
pks = [column.name for column in self.columns if column.is_pk]
|
||||
if not pks:
|
||||
pks = ["rowid"]
|
||||
|
|
@ -3124,7 +3143,7 @@ class Table(Queryable):
|
|||
if has_extracts:
|
||||
for i, key in enumerate(all_columns):
|
||||
if key in extracts:
|
||||
record_values[i] = self.db[extracts[key]].lookup(
|
||||
record_values[i] = self.db.table(extracts[key]).lookup(
|
||||
{"value": record_values[i]}
|
||||
)
|
||||
values.append(record_values)
|
||||
|
|
@ -3145,7 +3164,7 @@ class Table(Queryable):
|
|||
)
|
||||
if key in extracts:
|
||||
extract_table = extracts[key]
|
||||
value = self.db[extract_table].lookup({"value": value})
|
||||
value = self.db.table(extract_table).lookup({"value": value})
|
||||
record_values.append(value)
|
||||
values.append(record_values)
|
||||
|
||||
|
|
@ -3789,6 +3808,7 @@ class Table(Queryable):
|
|||
)
|
||||
)
|
||||
try:
|
||||
assert pk is not None, "pk cannot be None"
|
||||
return rows[0][pk]
|
||||
except IndexError:
|
||||
return self.insert(
|
||||
|
|
|
|||
|
|
@ -1,3 +1,6 @@
|
|||
import sqlite3
|
||||
|
||||
import click
|
||||
from pluggy import HookimplMarker
|
||||
from pluggy import HookspecMarker
|
||||
|
||||
|
|
@ -6,10 +9,10 @@ hookimpl = HookimplMarker("sqlite_utils")
|
|||
|
||||
|
||||
@hookspec
|
||||
def register_commands(cli):
|
||||
def register_commands(cli: click.Group) -> None:
|
||||
"""Register additional CLI commands, e.g. 'sqlite-utils mycommand ...'"""
|
||||
|
||||
|
||||
@hookspec
|
||||
def prepare_connection(conn):
|
||||
def prepare_connection(conn: sqlite3.Connection) -> None:
|
||||
"""Modify SQLite connection in some way e.g. register custom SQL functions"""
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
from typing import Any, Dict, List
|
||||
|
||||
import pluggy
|
||||
import sys
|
||||
from . import hookspecs
|
||||
|
||||
pm = pluggy.PluginManager("sqlite_utils")
|
||||
pm: pluggy.PluginManager = pluggy.PluginManager("sqlite_utils")
|
||||
pm.add_hookspecs(hookspecs)
|
||||
|
||||
if not getattr(sys, "_called_from_test", False):
|
||||
|
|
@ -10,13 +12,14 @@ if not getattr(sys, "_called_from_test", False):
|
|||
pm.load_setuptools_entrypoints("sqlite_utils")
|
||||
|
||||
|
||||
def get_plugins():
|
||||
plugins = []
|
||||
def get_plugins() -> List[Dict[str, Any]]:
|
||||
plugins: List[Dict[str, Any]] = []
|
||||
plugin_to_distinfo = dict(pm.list_plugin_distinfo())
|
||||
for plugin in pm.get_plugins():
|
||||
plugin_info = {
|
||||
hookcallers = pm.get_hookcallers(plugin)
|
||||
plugin_info: Dict[str, Any] = {
|
||||
"name": plugin.__name__,
|
||||
"hooks": [h.name for h in pm.get_hookcallers(plugin)],
|
||||
"hooks": [h.name for h in hookcallers] if hookcallers else [],
|
||||
}
|
||||
distinfo = plugin_to_distinfo.get(plugin)
|
||||
if distinfo:
|
||||
|
|
|
|||
|
|
@ -1,11 +1,23 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
from dateutil import parser
|
||||
import json
|
||||
|
||||
IGNORE = object()
|
||||
SET_NULL = object()
|
||||
IGNORE: object = object()
|
||||
SET_NULL: object = object()
|
||||
|
||||
|
||||
def parsedate(value, dayfirst=False, yearfirst=False, errors=None):
|
||||
def parsedate(
|
||||
value: Optional[str],
|
||||
dayfirst: bool = False,
|
||||
yearfirst: bool = False,
|
||||
errors: Optional[object] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Parse a date and convert it to ISO date format: yyyy-mm-dd
|
||||
\b
|
||||
|
|
@ -31,7 +43,12 @@ def parsedate(value, dayfirst=False, yearfirst=False, errors=None):
|
|||
raise
|
||||
|
||||
|
||||
def parsedatetime(value, dayfirst=False, yearfirst=False, errors=None):
|
||||
def parsedatetime(
|
||||
value: Optional[str],
|
||||
dayfirst: bool = False,
|
||||
yearfirst: bool = False,
|
||||
errors: Optional[object] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Parse a datetime and convert it to ISO datetime format: yyyy-mm-ddTHH:MM:SS
|
||||
\b
|
||||
|
|
@ -53,7 +70,9 @@ def parsedatetime(value, dayfirst=False, yearfirst=False, errors=None):
|
|||
raise
|
||||
|
||||
|
||||
def jsonsplit(value, delimiter=",", type=str):
|
||||
def jsonsplit(
|
||||
value: str, delimiter: str = ",", type: Callable[[str], Any] = str
|
||||
) -> str:
|
||||
"""
|
||||
Convert a string like a,b,c into a JSON array ["a", "b", "c"]
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -8,7 +8,22 @@ import itertools
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import Dict, cast, BinaryIO, Iterable, Iterator, Optional, Tuple, Type
|
||||
from typing import (
|
||||
Any,
|
||||
BinaryIO,
|
||||
Callable,
|
||||
Dict,
|
||||
Generator,
|
||||
Iterable,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
Set,
|
||||
Tuple,
|
||||
Type,
|
||||
Union,
|
||||
cast,
|
||||
)
|
||||
|
||||
import click
|
||||
|
||||
|
|
@ -44,17 +59,19 @@ SPATIALITE_PATHS = (
|
|||
ORIGINAL_CSV_FIELD_SIZE_LIMIT = csv.field_size_limit()
|
||||
|
||||
|
||||
class _CloseableIterator(Iterator[dict]):
|
||||
class _CloseableIterator(Iterator[Dict[str, Any]]):
|
||||
"""Iterator wrapper that closes a file when iteration is complete."""
|
||||
|
||||
def __init__(self, iterator: Iterator[dict], closeable: io.IOBase):
|
||||
def __init__(
|
||||
self, iterator: Iterator[Dict[str, Any]], closeable: io.IOBase
|
||||
) -> None:
|
||||
self._iterator = iterator
|
||||
self._closeable = closeable
|
||||
|
||||
def __iter__(self) -> "_CloseableIterator":
|
||||
return self
|
||||
|
||||
def __next__(self) -> dict:
|
||||
def __next__(self) -> Dict[str, Any]:
|
||||
try:
|
||||
return next(self._iterator)
|
||||
except StopIteration:
|
||||
|
|
@ -65,7 +82,7 @@ class _CloseableIterator(Iterator[dict]):
|
|||
self._closeable.close()
|
||||
|
||||
|
||||
def maximize_csv_field_size_limit():
|
||||
def maximize_csv_field_size_limit() -> None:
|
||||
"""
|
||||
Increase the CSV field size limit to the maximum possible.
|
||||
"""
|
||||
|
|
@ -108,20 +125,25 @@ def find_spatialite() -> Optional[str]:
|
|||
return None
|
||||
|
||||
|
||||
def suggest_column_types(records):
|
||||
all_column_types = {}
|
||||
def suggest_column_types(
|
||||
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))
|
||||
return types_for_column_types(all_column_types)
|
||||
|
||||
|
||||
def types_for_column_types(all_column_types):
|
||||
column_types = {}
|
||||
def types_for_column_types(
|
||||
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:
|
||||
types.discard(None.__class__)
|
||||
t: type
|
||||
if {None.__class__} == types:
|
||||
t = str
|
||||
elif len(types) == 1:
|
||||
|
|
@ -143,7 +165,7 @@ def types_for_column_types(all_column_types):
|
|||
return column_types
|
||||
|
||||
|
||||
def column_affinity(column_type):
|
||||
def column_affinity(column_type: str) -> type:
|
||||
# Implementation of SQLite affinity rules from
|
||||
# https://www.sqlite.org/datatype3.html#determination_of_column_affinity
|
||||
assert isinstance(column_type, str)
|
||||
|
|
@ -162,7 +184,7 @@ def column_affinity(column_type):
|
|||
return float
|
||||
|
||||
|
||||
def decode_base64_values(doc):
|
||||
def decode_base64_values(doc: Dict[str, Any]) -> Dict[str, Any]:
|
||||
# Looks for '{"$base64": true..., "encoded": ...}' values and decodes them
|
||||
to_fix = [
|
||||
k
|
||||
|
|
@ -177,23 +199,25 @@ def decode_base64_values(doc):
|
|||
|
||||
|
||||
class UpdateWrapper:
|
||||
def __init__(self, wrapped, update):
|
||||
def __init__(self, wrapped: Any, update: Callable[[int], None]) -> None:
|
||||
self._wrapped = wrapped
|
||||
self._update = update
|
||||
|
||||
def __iter__(self):
|
||||
def __iter__(self) -> Iterator[Any]:
|
||||
for line in self._wrapped:
|
||||
self._update(len(line))
|
||||
yield line
|
||||
|
||||
def read(self, size=-1):
|
||||
def read(self, size: int = -1) -> Any:
|
||||
data = self._wrapped.read(size)
|
||||
self._update(len(data))
|
||||
return data
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def file_progress(file, silent=False, **kwargs):
|
||||
def file_progress(
|
||||
file: Any, silent: bool = False, **kwargs: Any
|
||||
) -> Generator[Any, None, None]:
|
||||
if silent:
|
||||
yield file
|
||||
return
|
||||
|
|
@ -231,28 +255,30 @@ class RowError(Exception):
|
|||
|
||||
|
||||
def _extra_key_strategy(
|
||||
reader: Iterable[dict],
|
||||
reader: Iterable[Dict[Any, Any]],
|
||||
ignore_extras: Optional[bool] = False,
|
||||
extras_key: Optional[str] = None,
|
||||
) -> Iterable[dict]:
|
||||
) -> Iterable[Dict[str, Any]]:
|
||||
# Logic for handling CSV rows with more values than there are headings
|
||||
for row in reader:
|
||||
# DictReader adds a 'None' key with extra row values
|
||||
if None not in row:
|
||||
yield row
|
||||
row_with_optional_none = cast(Dict[Optional[str], Any], row)
|
||||
if None not in row_with_optional_none:
|
||||
yield cast(Dict[str, Any], row)
|
||||
elif ignore_extras:
|
||||
# ignoring row.pop(none) because of this issue:
|
||||
# https://github.com/simonw/sqlite-utils/issues/440#issuecomment-1155358637
|
||||
row.pop(None) # type: ignore
|
||||
yield row
|
||||
row_with_optional_none.pop(None)
|
||||
yield cast(Dict[str, Any], row)
|
||||
elif not extras_key:
|
||||
extras = row.pop(None) # type: ignore
|
||||
extras = row_with_optional_none.pop(None)
|
||||
raise RowError(
|
||||
"Row {} contained these extra values: {}".format(row, extras)
|
||||
)
|
||||
else:
|
||||
row[extras_key] = row.pop(None) # type: ignore
|
||||
yield row
|
||||
row_out = cast(Dict[str, Any], row)
|
||||
row_out[extras_key] = row_with_optional_none.pop(None)
|
||||
yield row_out
|
||||
|
||||
|
||||
def rows_from_file(
|
||||
|
|
@ -262,7 +288,7 @@ def rows_from_file(
|
|||
encoding: Optional[str] = None,
|
||||
ignore_extras: Optional[bool] = False,
|
||||
extras_key: Optional[str] = None,
|
||||
) -> Tuple[Iterable[dict], Format]:
|
||||
) -> Tuple[Iterable[Dict[str, Any]], Format]:
|
||||
"""
|
||||
Load a sequence of dictionaries from a file-like object containing one of four different formats.
|
||||
|
||||
|
|
@ -376,10 +402,10 @@ class TypeTracker:
|
|||
db["creatures"].transform(types=tracker.types)
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.trackers = {}
|
||||
def __init__(self) -> None:
|
||||
self.trackers: Dict[str, "ValueTracker"] = {}
|
||||
|
||||
def wrap(self, iterator: Iterable[dict]) -> Iterable[dict]:
|
||||
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.
|
||||
|
|
@ -402,25 +428,27 @@ class TypeTracker:
|
|||
|
||||
|
||||
class ValueTracker:
|
||||
def __init__(self):
|
||||
couldbe: Dict[str, Callable[[Any], bool]]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.couldbe = {key: getattr(self, "test_" + key) for key in self.get_tests()}
|
||||
|
||||
@classmethod
|
||||
def get_tests(cls):
|
||||
def get_tests(cls) -> List[str]:
|
||||
return [
|
||||
key.split("test_")[-1]
|
||||
for key in cls.__dict__.keys()
|
||||
if key.startswith("test_")
|
||||
]
|
||||
|
||||
def test_integer(self, value):
|
||||
def test_integer(self, value: Any) -> bool:
|
||||
try:
|
||||
int(value)
|
||||
return True
|
||||
except (ValueError, TypeError):
|
||||
return False
|
||||
|
||||
def test_float(self, value):
|
||||
def test_float(self, value: Any) -> bool:
|
||||
try:
|
||||
float(value)
|
||||
return True
|
||||
|
|
@ -431,7 +459,7 @@ class ValueTracker:
|
|||
return self.guessed_type + ": possibilities = " + repr(self.couldbe)
|
||||
|
||||
@property
|
||||
def guessed_type(self):
|
||||
def guessed_type(self) -> str:
|
||||
options = set(self.couldbe.keys())
|
||||
# Return based on precedence
|
||||
for key in self.get_tests():
|
||||
|
|
@ -439,10 +467,10 @@ class ValueTracker:
|
|||
return key
|
||||
return "text"
|
||||
|
||||
def evaluate(self, value):
|
||||
def evaluate(self, value: Any) -> None:
|
||||
if not value or not self.couldbe:
|
||||
return
|
||||
not_these = []
|
||||
not_these: List[str] = []
|
||||
for name, test in self.couldbe.items():
|
||||
if not test(value):
|
||||
not_these.append(name)
|
||||
|
|
@ -451,18 +479,18 @@ class ValueTracker:
|
|||
|
||||
|
||||
class NullProgressBar:
|
||||
def __init__(self, *args):
|
||||
def __init__(self, *args: Any) -> None:
|
||||
self.args = args
|
||||
|
||||
def __iter__(self):
|
||||
def __iter__(self) -> Iterator[Any]:
|
||||
yield from self.args[0]
|
||||
|
||||
def update(self, value):
|
||||
def update(self, value: int) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def progressbar(*args, **kwargs):
|
||||
def progressbar(*args: Any, **kwargs: Any) -> Generator[Any, None, None]:
|
||||
silent = kwargs.pop("silent")
|
||||
if silent:
|
||||
yield NullProgressBar(*args)
|
||||
|
|
@ -471,25 +499,27 @@ def progressbar(*args, **kwargs):
|
|||
yield bar
|
||||
|
||||
|
||||
def _compile_code(code, imports, variable="value"):
|
||||
globals = {"r": recipes, "recipes": recipes}
|
||||
def _compile_code(
|
||||
code: str, imports: Iterable[str], variable: str = "value"
|
||||
) -> Callable[..., Any]:
|
||||
globals_dict: Dict[str, Any] = {"r": recipes, "recipes": recipes}
|
||||
# Handle imports first so they're available for all approaches
|
||||
for import_ in imports:
|
||||
globals[import_.split(".")[0]] = __import__(import_)
|
||||
globals_dict[import_.split(".")[0]] = __import__(import_)
|
||||
|
||||
# If user defined a convert() function, return that
|
||||
try:
|
||||
exec(code, globals)
|
||||
return globals["convert"]
|
||||
exec(code, globals_dict)
|
||||
return cast(Callable[..., Any], globals_dict["convert"])
|
||||
except (AttributeError, SyntaxError, NameError, KeyError, TypeError):
|
||||
pass
|
||||
|
||||
# Check if code is a direct callable reference
|
||||
# e.g. "r.parsedate" instead of "r.parsedate(value)"
|
||||
try:
|
||||
fn = eval(code, globals)
|
||||
fn = eval(code, globals_dict)
|
||||
if callable(fn):
|
||||
return fn
|
||||
return cast(Callable[..., Any], fn)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
|
@ -514,11 +544,11 @@ def _compile_code(code, imports, variable="value"):
|
|||
if code_o is None:
|
||||
raise SyntaxError("Could not compile code")
|
||||
|
||||
exec(code_o, globals)
|
||||
return globals["fn"]
|
||||
exec(code_o, globals_dict)
|
||||
return cast(Callable[..., Any], globals_dict["fn"])
|
||||
|
||||
|
||||
def chunks(sequence: Iterable, size: int) -> Iterable[Iterable]:
|
||||
def chunks(sequence: Iterable[Any], size: int) -> Iterable[Iterable[Any]]:
|
||||
"""
|
||||
Iterate over chunks of the sequence of the given size.
|
||||
|
||||
|
|
@ -530,7 +560,9 @@ def chunks(sequence: Iterable, size: int) -> Iterable[Iterable]:
|
|||
yield itertools.chain([item], itertools.islice(iterator, size - 1))
|
||||
|
||||
|
||||
def hash_record(record: Dict, keys: Optional[Iterable[str]] = None):
|
||||
def hash_record(
|
||||
record: Dict[str, Any], keys: Optional[Iterable[str]] = None
|
||||
) -> str:
|
||||
"""
|
||||
``record`` should be a Python dictionary. Returns a sha1 hash of the
|
||||
keys and values in that record.
|
||||
|
|
@ -551,7 +583,7 @@ def hash_record(record: Dict, 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 = record
|
||||
to_hash: Dict[str, Any] = record
|
||||
if keys is not None:
|
||||
to_hash = {key: record[key] for key in keys}
|
||||
return hashlib.sha1(
|
||||
|
|
@ -561,7 +593,7 @@ def hash_record(record: Dict, keys: Optional[Iterable[str]] = None):
|
|||
).hexdigest()
|
||||
|
||||
|
||||
def _flatten(d):
|
||||
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):
|
||||
|
|
@ -570,7 +602,7 @@ def _flatten(d):
|
|||
yield key, value
|
||||
|
||||
|
||||
def flatten(row: dict) -> dict:
|
||||
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}``
|
||||
|
||||
|
|
|
|||
|
|
@ -41,9 +41,9 @@ def test_convert_help():
|
|||
result = CliRunner().invoke(cli.cli, ["convert", "--help"])
|
||||
assert result.exit_code == 0
|
||||
for expected in (
|
||||
"r.jsonsplit(value, ",
|
||||
"r.parsedate(value, ",
|
||||
"r.parsedatetime(value, ",
|
||||
"r.jsonsplit(value:",
|
||||
"r.parsedate(value:",
|
||||
"r.parsedatetime(value:",
|
||||
):
|
||||
assert expected in result.output
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue