thds.tabularasa 0.13.0__py3-none-any.whl

thds/tabularasa/schema/compilation/attrs.py

@@ -0,0 +1,257 @@
+import typing as ty
+from operator import itemgetter
+from textwrap import wrap
+
+import typing_extensions
+
+import thds.tabularasa.loaders.util
+from thds.tabularasa.schema import metaschema
+
+from ._format import autoformat
+from .util import (
+    AUTOGEN_DISCLAIMER,
+    VarName,
+    _indent,
+    _wrap_lines_with_prefix,
+    constructor_template,
+    render_blob_store_def,
+    render_constructor,
+    sorted_class_names_for_import,
+)
+
+DEFAULT_LINE_WIDTH = 88
+
+REMOTE_BLOB_STORE_VAR_NAME = "REMOTE_BLOB_STORE"
+
+DOCSTRING_PARAM_TEMPLATE = """:param {name}: {doc}"""
+
+CUSTOM_TYPE_DEF_TEMPLATE = """{comment}{name} = {type}"""
+
+ATTRS_CLASS_DEF_TEMPLATE = """@attr.s(auto_attribs=True, frozen=True)
+class {class_name}:
+    \"\"\"{doc}
+
+    {params}
+    \"\"\"
+
+    {fields}
+"""
+
+ATTRS_FIELD_DEF_TEMPLATE_BASIC = "{name}: {type}"
+
+ATTRS_LOADER_TEMPLATE = constructor_template(
+    thds.tabularasa.loaders.util.AttrsParquetLoader,
+    exclude=["filename"],
+    type_params=["{record_type}"],
+)
+
+
+def render_type_def(
+    type_: metaschema.CustomType,
+    build_options: metaschema.BuildOptions,
+) -> str:
+    type_literal = type_.python_type_def_literal(build_options)
+
+    if build_options.type_constraint_comments:
+        comment = type_.comment
+        if comment:
+            lines = wrap(comment, DEFAULT_LINE_WIDTH - 2)
+            comment = "\n".join("# " + line for line in lines) + "\n"
+        else:
+            comment = ""
+    else:
+        comment = ""
+
+    return CUSTOM_TYPE_DEF_TEMPLATE.format(comment=comment, name=type_.class_name, type=type_literal)
+
+
+def render_attr_field_def(
+    column: metaschema.Column, build_options: metaschema.BuildOptions, builtin: bool = False
+) -> str:
+    type_literal = column.python_type_literal(build_options=build_options, builtin=builtin)
+    return ATTRS_FIELD_DEF_TEMPLATE_BASIC.format(name=column.snake_case_name, type=type_literal)
+
+
+def render_attrs_table_schema(table: metaschema.Table, build_options: metaschema.BuildOptions) -> str:
+    field_defs = []
+    params = []
+
+    for column in table.columns:
+        field_def = render_attr_field_def(column, builtin=False, build_options=build_options)
+        field_defs.append(field_def)
+        doc = _wrap_lines_with_prefix(
+            column.doc,
+            DEFAULT_LINE_WIDTH - 4,
+            first_line_prefix_len=len(f":param {column.snake_case_name}: "),
+            trailing_line_indent=2,
+        )
+        params.append(
+            DOCSTRING_PARAM_TEMPLATE.format(
+                name=column.snake_case_name,
+                doc=doc,
+            )
+        )
+
+    table_doc = _wrap_lines_with_prefix(
+        table.doc,
+        DEFAULT_LINE_WIDTH - 4,
+        first_line_prefix_len=3,  # triple quotes
+        trailing_line_indent=0,
+    )
+
+    return ATTRS_CLASS_DEF_TEMPLATE.format(
+        class_name=table.class_name,
+        doc=_indent(table_doc),
+        params=_indent("\n".join(params)),
+        fields=_indent("\n".join(field_defs)),
+    )
+
+
+PYARROW_SCHEMAS_QUALIFIED_IMPORT = "pyarrow_schemas"
+
+
+class ImportsAndCode(ty.NamedTuple):
+    """Couples code and its required imports."""
+
+    third_party_imports: ty.List[str]
+    tabularasa_imports: ty.List[str]
+    code: ty.List[str]
+
+
+def render_attrs_loaders(
+    schema: metaschema.Schema,
+    package: str,
+) -> ImportsAndCode:
+    data_dir = schema.build_options.package_data_dir
+    render_pyarrow_schemas = schema.build_options.pyarrow
+    import_lines = list()
+    if render_pyarrow_schemas:
+        import_lines.append("\n")
+        import_lines.append(f"from . import pyarrow as {PYARROW_SCHEMAS_QUALIFIED_IMPORT}")
+
+    return ImportsAndCode(
+        list(),
+        import_lines,
+        [
+            render_constructor(
+                ATTRS_LOADER_TEMPLATE,
+                kwargs=dict(
+                    record_type=VarName(table.class_name),
+                    table_name=table.snake_case_name,
+                    type_=VarName(table.class_name),
+                    package=package,
+                    data_dir=data_dir,
+                    md5=table.md5,
+                    blob_store=(
+                        None
+                        if schema.remote_blob_store is None or table.md5 is None
+                        else VarName(REMOTE_BLOB_STORE_VAR_NAME)
+                    ),
+                    pyarrow_schema=(
+                        VarName(
+                            f"{PYARROW_SCHEMAS_QUALIFIED_IMPORT}.{table.snake_case_name}_pyarrow_schema"
+                        )
+                        if render_pyarrow_schemas
+                        else None
+                    ),
+                ),
+                var_name=f"load_{table.snake_case_name}",
+            )
+            for table in schema.package_tables
+        ],
+    )
+
+
+def render_attrs_type_defs(
+    schema: metaschema.Schema,
+) -> ImportsAndCode:
+    # custom types
+    defined_custom_types = schema.defined_types
+    type_defs = [
+        render_type_def(
+            type_,
+            build_options=schema.build_options,
+        )
+        for type_ in sorted(defined_custom_types, key=lambda type_: type_.name)
+    ]
+
+    import_lines = list()
+    # external type imports
+    sep = ",\n    "
+    for module_name, class_names in sorted(schema.external_type_imports.items(), key=itemgetter(0)):
+        import_lines.append(
+            f"from {module_name} import (\n    {sep.join(sorted_class_names_for_import(class_names))},\n)\n"
+        )
+
+    return ImportsAndCode([], import_lines, type_defs)
+
+
+def _render_attrs_schema(
+    schema: metaschema.Schema,
+    type_defs: ImportsAndCode,
+    loader_defs: ty.Optional[ImportsAndCode],
+) -> str:
+    loader_defs = loader_defs or ImportsAndCode([], [], [])
+    assert loader_defs, "Loaders are optional but the line above is not"
+
+    # attrs record types
+    table_defs = [
+        render_attrs_table_schema(table, schema.build_options) for table in schema.package_tables
+    ]
+
+    # imports
+    stdlib_imports = sorted(schema.attrs_required_imports)
+    import_extensions = typing_extensions.__name__ in stdlib_imports
+    if import_extensions:
+        stdlib_imports.remove(typing_extensions.__name__)
+    import_lines = [f"import {module}\n" for module in stdlib_imports]
+
+    if import_lines:
+        import_lines.append("\n")
+    import_lines.append("import attr\n")
+    if import_extensions:
+        import_lines.append(f"import {typing_extensions.__name__}\n")
+
+    import_lines.append("\n")
+
+    if loader_defs.code:
+        import_lines.append(f"import {thds.tabularasa.loaders.util.__name__}\n")
+    if schema.remote_blob_store is not None:
+        import_lines.append(f"import {thds.tabularasa.schema.files.__name__}\n")
+
+    import_lines.extend(type_defs.tabularasa_imports)
+    import_lines.extend(loader_defs.tabularasa_imports)
+
+    # globals
+    global_var_defs = []
+    if schema.remote_blob_store is not None:
+        global_var_defs.append(
+            render_blob_store_def(schema.remote_blob_store, REMOTE_BLOB_STORE_VAR_NAME)
+        )
+
+    imports = "".join(import_lines)
+    globals_ = "\n".join(global_var_defs)
+    types = "\n".join(type_defs.code)
+    classes = "\n\n".join(table_defs)
+    loaders = "\n\n".join(loader_defs.code)
+
+    # module
+    return autoformat(
+        f"{imports}\n# {AUTOGEN_DISCLAIMER}\n\n{globals_}\n\n{types}\n\n\n{classes}\n\n{loaders}\n"
+    )
+
+
+def render_attrs_module(
+    schema: metaschema.Schema,
+    package: str,
+    loader_defs: ty.Optional[ImportsAndCode] = None,
+) -> str:
+    if loader_defs is None:
+        loader_defs = (
+            render_attrs_loaders(schema, package) if schema.build_options.package_data_dir else None
+        )
+    return _render_attrs_schema(
+        schema,
+        render_attrs_type_defs(schema),
+        loader_defs,
+    )

thds/tabularasa/schema/compilation/attrs_sqlite.py

@@ -0,0 +1,278 @@
+from itertools import chain
+from logging import getLogger
+from typing import Dict, List, Optional, Tuple
+
+import thds.tabularasa.loaders.util
+import thds.tabularasa.schema
+from thds.tabularasa.schema import metaschema
+from thds.tabularasa.schema.compilation.attrs import render_attr_field_def
+
+from ._format import autoformat
+from .sqlite import index_name
+from .util import AUTOGEN_DISCLAIMER, sorted_class_names_for_import
+
+_LOGGER = getLogger(__name__)
+
+PACKAGE_VARNAME = "PACKAGE"
+DB_PATH_VARNAME = "DB_PATH"
+
+ATTRS_MODULE_NAME = ".attrs"
+
+LINE_WIDTH = 88
+
+COLUMN_LINESEP = ",\n                  "
+
+ATTRS_CLASS_LOADER_TEMPLATE = """class {class_name}Loader:
+
+    def __init__(self, db: util.%s):
+        self._db = db
+        self._record = util.%s({class_name})
+
+    {accessors}
+""" % (
+    thds.tabularasa.loaders.util.AttrsSQLiteDatabase.__name__,
+    thds.tabularasa.loaders.sqlite_util.sqlite_constructor_for_record_type.__name__,  # type: ignore
+)
+
+ATTRS_INDEX_ACCESSOR_TEMPLATE = """
+    def {method_name}(self, {typed_args}) -> typing.{return_type}[{class_name}]:
+        return self._db.sqlite_{index_kind}_query(
+            self._record,
+            \"\"\"
+            SELECT
+              {columns}
+            FROM {table_name}
+            INDEXED BY {index_name}
+            WHERE {condition};
+            \"\"\",
+            ({args},),
+        )
+"""
+
+ATTRS_BULK_INDEX_ACCESSOR_TEMPLATE = """
+    def {method_name}_bulk(self, {arg_name}: typing.Collection[{typed_args}]) -> typing.{return_type}[{class_name}]:
+        if {arg_name}:
+            return self._db.sqlite_{index_kind}_query(
+                self._record,
+                f\"\"\"
+                SELECT
+                  {columns}
+                FROM {table_name}
+                INDEXED BY {index_name}
+                WHERE {condition};
+                \"\"\",
+                {arg_name},
+                single_col={single_col},
+            )
+        else:
+            return iter(())
+"""
+
+ATTRS_MAIN_LOADER_TEMPLATE = """class SQLiteLoader:
+    def __init__(
+        self,
+        package: typing.Optional[str] = %s,
+        db_path: str = %s,
+        cache_size: int = util.DEFAULT_ATTR_SQLITE_CACHE_SIZE,
+        mmap_size: int = util.DEFAULT_MMAP_BYTES,
+    ):
+        self._db = util.%s(package=package, db_path=db_path, cache_size=cache_size, mmap_size=mmap_size)
+        {table_loaders}
+""" % (
+    PACKAGE_VARNAME,
+    DB_PATH_VARNAME,
+    thds.tabularasa.loaders.util.AttrsSQLiteDatabase.__name__,
+)
+
+
+def render_attrs_loader_schema(table: metaschema.Table, build_options: metaschema.BuildOptions) -> str:
+    accessor_defs = []
+    column_lookup = {col.name: col for col in table.columns}
+    unq_constraints = {frozenset(c.unique) for c in table.unique_constraints}
+
+    if table.primary_key:
+        accessor_defs.append(
+            render_accessor_method(
+                table, table.primary_key, column_lookup, pk=True, build_options=build_options
+            )
+        )
+        accessor_defs.append(
+            render_accessor_method(
+                table, table.primary_key, column_lookup, pk=True, bulk=True, build_options=build_options
+            )
+        )
+
+    for idx in table.indexes:
+        unique = frozenset(idx) in unq_constraints
+        accessor_defs.append(
+            render_accessor_method(
+                table, idx, column_lookup, pk=False, unique=unique, build_options=build_options
+            )
+        )
+        accessor_defs.append(
+            render_accessor_method(
+                table,
+                idx,
+                column_lookup,
+                pk=False,
+                unique=unique,
+                bulk=True,
+                build_options=build_options,
+            )
+        )
+
+    accessors = "".join(accessor_defs).strip()
+    return ATTRS_CLASS_LOADER_TEMPLATE.format(
+        class_name=table.class_name,
+        accessors=accessors,
+    )
+
+
+def render_accessor_method(
+    table: metaschema.Table,
+    index_columns: Tuple[str, ...],
+    column_lookup: Dict[str, metaschema.Column],
+    build_options: metaschema.BuildOptions,
+    pk: bool = False,
+    unique: bool = False,
+    bulk: bool = False,
+) -> str:
+    index_column_names = tuple(map(metaschema.snake_case, index_columns))
+    method_name = "pk" if pk else f"idx_{'_'.join(index_column_names)}"
+    index_kind = "bulk" if bulk else ("pk" if pk or unique else "index")
+    return_type = "Iterator" if bulk else ("Optional" if pk or unique else "List")
+    # we use the `IS` operator to allow for comparison in case of nullable index columns
+    arg_name = "__".join(index_column_names)
+    nullsafe_condition = " AND ".join([f"{column} IS (?)" for column in index_column_names])
+    if bulk:
+        columns = [column_lookup[col] for col in index_columns]
+        has_null_cols = any(col.nullable for col in columns)
+        if len(index_column_names) == 1:
+            # for single-column indexes, we need to unpack the single value from the tuple
+            typed_args = columns[0].python_type_literal(build_options=build_options, builtin=True)
+            if has_null_cols:
+                # sqlite IN operator unfortunately doesn't support a NULL == NULL variant, the way that IS does for =
+                condition = f"{{' OR '.join(['{index_column_names[0]} IS (?)'] * len({arg_name}))}}"
+            else:
+                condition = f"{index_column_names[0]} IN ({{','.join('?' * len({arg_name}))}})"
+            single_col = "True"
+        else:
+            type_strs = ", ".join(
+                col.python_type_literal(build_options=build_options, builtin=True) for col in columns
+            )
+            typed_args = f"typing.Tuple[{type_strs}]"
+            if has_null_cols:
+                # sqlite IN operator unfortunately doesn't support a NULL == NULL variant, the way that IS does for =
+                param_tuple = f"({nullsafe_condition})"
+                condition = f"{{' OR '.join(['{param_tuple}'] * len({arg_name}))}}"
+            else:
+                param_tuple = f"({', '.join('?' * len(index_column_names))})"
+                condition = f"({', '.join(index_column_names)}) IN ({{','.join(['{param_tuple}'] * len({arg_name}))}})"
+            single_col = "False"
+    else:
+        condition = nullsafe_condition
+        typed_args = ", ".join(
+            [
+                render_attr_field_def(column_lookup[col], builtin=True, build_options=build_options)
+                for col in index_columns
+            ]
+        )
+        single_col = ""
+
+    return (ATTRS_BULK_INDEX_ACCESSOR_TEMPLATE if bulk else ATTRS_INDEX_ACCESSOR_TEMPLATE).format(
+        class_name=table.class_name,
+        method_name=method_name,
+        # use builtin types (as opposed to e.g. Literals and Newtypes) to make the API simpler to use
+        typed_args=typed_args,
+        arg_name=arg_name,
+        single_col=single_col,
+        return_type=return_type,
+        index_kind=index_kind,
+        table_name=table.snake_case_name,
+        columns=COLUMN_LINESEP.join(c.snake_case_name for c in table.columns),
+        index_name=index_name(table.snake_case_name, *index_column_names),
+        condition=condition,
+        args=", ".join(index_column_names),
+    )
+
+
+def render_attrs_main_loader(
+    tables: List[metaschema.Table],
+) -> str:
+    loader_instance_defs = [
+        f"self.{table.snake_case_name} = {table.class_name}Loader(self._db)" for table in tables
+    ]
+    table_loaders = "\n        ".join(loader_instance_defs)
+    return ATTRS_MAIN_LOADER_TEMPLATE.format(table_loaders=table_loaders)
+
+
+def _import_lines(
+    tables: List[metaschema.Table],
+    attrs_module_name: Optional[str],
+    build_options: metaschema.BuildOptions,
+):
+    # need typing always for List and Optional
+    stdlib_imports = sorted(
+        {
+            "typing",
+            *chain.from_iterable(t.attrs_sqlite_required_imports(build_options) for t in tables),
+        }
+    )
+    import_lines = [f"import {module}\n" for module in stdlib_imports]
+    if import_lines:
+        import_lines.append("\n")
+
+    import_lines.append(f"import {thds.tabularasa.loaders.sqlite_util.__name__} as util\n")
+    import_lines.append("\n")
+    if attrs_module_name is not None:
+        import_lines.append(f"from {attrs_module_name} import (\n")
+        attrs_module_classnames = {table.class_name for table in tables}
+        import_lines.extend(
+            f"    {name},\n" for name in sorted_class_names_for_import(attrs_module_classnames)
+        )
+        import_lines.append(")\n")
+
+    return "".join(import_lines)
+
+
+def _has_index(table: metaschema.Table) -> bool:
+    return (table.primary_key is not None) or len(table.indexes) > 0
+
+
+def render_attrs_sqlite_schema(
+    schema: metaschema.Schema,
+    package: str = "",
+    db_path: str = "",
+    attrs_module_name: Optional[str] = ATTRS_MODULE_NAME,
+) -> str:
+    has_database_loader = bool(package and db_path)
+
+    # do not generate a SQLite loader if there is no primary key or index defined on the table def
+    tables = [table for table in schema.package_tables if table.has_indexes]
+    tables_filtered = [table.name for table in schema.package_tables if not table.has_indexes]
+    if not tables:
+        _LOGGER.info(
+            f"Skipping SQLite loader generation for all tables: {tables_filtered}; none has any index "
+            f"specified"
+        )
+        return ""
+
+    if tables_filtered:
+        _LOGGER.info(
+            f"Skipping SQLite loader generation for the following "
+            f"tables because no indices or primary keys are defined: {', '.join(tables_filtered)}"
+        )
+
+    imports = _import_lines(tables, attrs_module_name, schema.build_options)
+    loader_defs = [render_attrs_loader_schema(table, schema.build_options) for table in tables]
+    loaders = "\n\n".join(loader_defs)
+
+    if has_database_loader:
+        constants = f'{PACKAGE_VARNAME} = "{package}"\n{DB_PATH_VARNAME} = "{db_path}"'
+        loader_def = render_attrs_main_loader(tables)
+    else:
+        constants = ""
+        loader_def = ""
+    return autoformat(
+        f"{imports}\n# {AUTOGEN_DISCLAIMER}\n\n{constants}\n\n\n{loaders}\n\n{loader_def}\n"
+    )

thds/tabularasa/schema/compilation/io.py

@@ -0,0 +1,96 @@
+import ast
+from functools import singledispatch
+from itertools import starmap, zip_longest
+from logging import getLogger
+from operator import itemgetter
+from pathlib import Path
+from typing import Any, Iterator, List, Tuple, Union
+
+_LOGGER = getLogger(__name__)
+
+AST_CODE_CONTEXT_VARS = {"lineno", "col_offset", "ctx", "end_lineno", "end_col_offset"}
+
+
+def ast_eq(ast1: ast.AST, ast2: ast.AST) -> bool:
+    """Return True if two python source strings are AST-equivalent, else False"""
+    return _ast_eq(ast1, ast2)
+
+
+def ast_vars(node: ast.AST) -> Iterator[Tuple[str, Any]]:
+    """Iterator of (name, value) tuples for all attributes of an AST node *except* for those that are
+    not abstract (e.g. line numbers and column offsets)"""
+    return (
+        (name, value)
+        for name, value in sorted(vars(node).items(), key=itemgetter(0))
+        if name not in AST_CODE_CONTEXT_VARS
+    )
+
+
+@singledispatch
+def _ast_eq(ast1: Any, ast2: Any) -> bool:
+    # base case, literal values (non-AST nodes)
+    return (type(ast1) is type(ast2)) and (ast1 == ast2)
+
+
+@_ast_eq.register(ast.AST)
+def _ast_eq_ast(ast1: ast.AST, ast2: ast.AST) -> bool:
+    if type(ast1) is not type(ast2):
+        return False
+    attrs1 = ast_vars(ast1)
+    attrs2 = ast_vars(ast2)
+    return all(
+        (name1 == name2) and _ast_eq(attr1, attr2)
+        for (name1, attr1), (name2, attr2) in zip(
+            attrs1,
+            attrs2,
+        )
+    )
+
+
+@_ast_eq.register(list)
+@_ast_eq.register(tuple)
+def _ast_eq_list(ast1: List[Any], ast2: List[Any]):
+    missing = object()
+    return all(starmap(_ast_eq, zip_longest(ast1, ast2, fillvalue=missing)))
+
+
+def write_if_ast_changed(source: str, path: Union[str, Path]):  # pragma: no cover
+    """Write the source code `source` to the file at `path`, but only if the file doesn't exist, or the
+    AST of the code therein differs from that of `source`"""
+    path = Path(path)
+    this_ast = ast.parse(source)
+
+    if path.exists():
+        with open(path, "r+") as f:
+            that_source = f.read()
+            try:
+                that_ast = ast.parse(that_source)
+            except SyntaxError:
+                _LOGGER.warning(
+                    f"syntax error in code at {path}; merge conflicts? Code will be overwritten"
+                )
+                rewrite = True
+                reason = "Invalid AST"
+            else:
+                rewrite = not ast_eq(this_ast, that_ast)
+                reason = "AST changed"
+
+            if rewrite:
+                _LOGGER.info(f"writing new generated code to {path}; {reason}")
+                f.seek(0)
+                f.truncate()
+                f.write(source)
+            else:
+                _LOGGER.info(f"leaving generated code at {path}; AST unchanged")
+    else:
+        _LOGGER.info(f"writing new generated code to {path}; no prior file")
+        with open(path, "w") as f:
+            f.write(source)
+
+
+def write_sql(source: str, path: Union[str, Path]):  # pragma: no cover
+    """Write the SQL source code `source` to the file at `path`"""
+    path = Path(path)
+    _LOGGER.info(f"writing new generated code to {path}; no prior file")
+    with open(path, "w") as f:
+        f.write(source)