thds.core 0.0.1__py3-none-any.whl → 1.31.20250116223856__py3-none-any.whl

thds/core/scope.py

@@ -0,0 +1,199 @@
+"""This allows the usage of ContextManagers that cover the entire body
+of a function without requiring invasive (and git-diff-increasing)
+"with" statements.
+
+Another way of looking at this is that it is essentially a
+decorator-driven `defer` in Go or `scope` in D. However, the semantics
+are slightly different in that the scope is not necessarily at the
+nearest function call boundary - we use Python's dynamic capabilities
+to look 'up' the stack until we find a scope that is usable, and then
+we embed the ContextManager in that scope.
+
+Generally, the usage will look something like this:
+
+```
+@scope.bound  # wrap a function with a scope that will exit when it returns
+def do_stuff(...):
+    foo = scope.enter(a_context_manager(...))  # enters the context manager via the nearest scope
+    # ...do some stuff
+    return bar
+    # context manager exits when nearest scope exits, which is right after function return.
+```
+
+where the traditional alternative would be:
+
+```
+def do_stuff(...):
+    with a_context_manager(...) as foo:
+        # ...do the same stuff
+        # ...but now your git diff is huge b/c you indented everything,
+        return bar
+    # context manager exits after `with` block closes, which is when the function returns.
+```
+
+Because we use ContextVar to perform the lookups of the nearest scope,
+this is actually pretty performant.  You do pay the cost of the
+wrapping function, which is higher than a `with` statement.
+
+"""
+
+import atexit
+import contextlib
+import inspect
+import sys
+import typing as ty
+from functools import wraps
+from logging import getLogger
+from uuid import uuid4
+
+from .inspect import get_caller_info
+from .stack_context import StackContext
+
+_KEYED_SCOPE_CONTEXTS: ty.Dict[str, StackContext[contextlib.ExitStack]] = dict()
+# all non-nil ExitStacks will be closed at application exit
+
+
+def _close_root_scopes_atexit():
+    for name, scope_sc in _KEYED_SCOPE_CONTEXTS.items():
+        scope = scope_sc()
+        if scope:
+            try:
+                scope.close()
+            except ValueError as ve:
+                print(f"Unable to close scope '{name}' at exit because {ve}", file=sys.stderr)
+
+
+atexit.register(_close_root_scopes_atexit)
+
+
+def _init_sc(key: str, val: contextlib.ExitStack):
+    """This should only ever be called at the root of/during import of a
+    module. It is _not_ threadsafe.
+    """
+    # normally you shouldn't create a StackContext except as a
+    # global.  in this case, we're dynamically storing _in_ a
+    # global dict, which is equivalent.
+    if key in _KEYED_SCOPE_CONTEXTS:
+        getLogger(__name__).warning(
+            f"Scope {key} already exists! If this is not importlib.reload, you have a problem."
+        )
+    _KEYED_SCOPE_CONTEXTS[key] = StackContext(key, val)
+
+
+F = ty.TypeVar("F", bound=ty.Callable)
+
+
+def _bound(key: str, func: F) -> F:
+    """A decorator that establishes a scope boundary for context managers
+    that can now be `enter`ed, and will then be exited when this
+    boundary is returned to.
+    """
+    if inspect.isgeneratorfunction(func):
+
+        @wraps(func)
+        def __scope_boundary_generator_wrap(*args, **kwargs):
+            if key not in _KEYED_SCOPE_CONTEXTS:
+                _init_sc(key, contextlib.ExitStack())  # this root stack will probably not get used
+
+            with _KEYED_SCOPE_CONTEXTS[key].set(contextlib.ExitStack()) as scoped_exit_stack:
+                with scoped_exit_stack:  # enter and exit the ExitStack itself
+                    ret = yield from func(*args, **kwargs)
+                    return ret  # weird syntax here, Python...
+
+        return ty.cast(F, __scope_boundary_generator_wrap)
+
+    if inspect.isasyncgenfunction(func):
+
+        @wraps(func)
+        async def __scope_boundary_async_generator_wrap(*args, **kwargs):
+            if key not in _KEYED_SCOPE_CONTEXTS:
+                _init_sc(key, contextlib.ExitStack())
+
+            with _KEYED_SCOPE_CONTEXTS[key].set(contextlib.ExitStack()) as scoped_exit_stack:
+                with scoped_exit_stack:
+                    async for ret in func(*args, **kwargs):
+                        yield ret
+
+        return ty.cast(F, __scope_boundary_async_generator_wrap)
+
+    if inspect.iscoroutinefunction(func):
+
+        @wraps(func)
+        async def __scope_boundary_coroutine_wrap(*args, **kwargs):
+            if key not in _KEYED_SCOPE_CONTEXTS:
+                _init_sc(key, contextlib.ExitStack())
+
+            with _KEYED_SCOPE_CONTEXTS[key].set(contextlib.ExitStack()) as scoped_exit_stack:
+                with scoped_exit_stack:
+                    return await func(*args, **kwargs)
+
+        return ty.cast(F, __scope_boundary_coroutine_wrap)
+
+    @wraps(func)
+    def __scope_boundary_wrap(*args, **kwargs):
+        if key not in _KEYED_SCOPE_CONTEXTS:
+            _init_sc(key, contextlib.ExitStack())  # this root stack will probably not get used
+
+        with _KEYED_SCOPE_CONTEXTS[key].set(contextlib.ExitStack()) as scoped_exit_stack:
+            with scoped_exit_stack:  # enter and exit the ExitStack itself
+                return func(*args, **kwargs)
+
+    return ty.cast(F, __scope_boundary_wrap)
+
+
+class NoScopeFound(Exception):
+    pass
+
+
+M = ty.TypeVar("M")
+
+
+def _enter(key: str, context: ty.ContextManager[M]) -> M:
+    """Call this to enter a ContextManager which will be exited at the
+    nearest scope boundary, without needing a with statement.
+    """
+    # this is fairly efficient - we don't walk up the stack; we simply
+    # use the stack-following ContextVar that was set up previously.
+    scope_context = _KEYED_SCOPE_CONTEXTS.get(key)
+    if scope_context:
+        return scope_context().enter_context(context)
+    raise NoScopeFound(f"No scope with the key {key} was found.")
+
+
+class Scope:
+    """Creating your own Scope isn't often necessary - often you just want
+    a basic scope around your function, so you can just use the default Scope,
+    which is created below.
+
+    However, in case it's important to your use case to be able to
+    have orthogonal scopes that can be entered further down the stack
+    and exited at a precise point further up, this makes it possible.
+
+    If you provide a key, it must be globally unique, and if it has
+    previously been created within the same application, an
+    AssertionError will be thrown. You do not need to provide a key.
+
+    These should be module-level/global objects under all
+    circumstances, as they share an internal global namespace.
+
+    """
+
+    def __init__(self, key: str = ""):
+        caller_info = get_caller_info(skip=1)
+        self.key = caller_info.module + "+" + (key or uuid4().hex)
+        _init_sc(self.key, contextlib.ExitStack())  # add root boundary
+
+    def bound(self, func: F) -> F:
+        """Add a boundary to this function which will close all of the
+        Contexts subsequently entered at the time this function exits.
+        """
+        return _bound(self.key, func)
+
+    def enter(self, context: ty.ContextManager[M]) -> M:
+        """Enter the provided Context with a future exit at the nearest boundary for this Scope."""
+        return _enter(self.key, context)
+
+
+default = Scope("__default_scope_stack")
+bound = default.bound
+enter = default.enter

thds/core/source.py

@@ -0,0 +1,238 @@
+"""Wrap openable, read-only data that is either locally-present or downloadable,
+
+yet will not be downloaded (if non-local) until it is actually opened or unwrapped.
+"""
+
+import os
+import typing as ty
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+
+from . import log
+from .files import is_file_uri, path_from_uri, to_uri
+from .hash_cache import filehash
+from .hashing import Hash
+from .types import StrOrPath
+
+
+class Downloader(ty.Protocol):
+    def __call__(self, hash: ty.Optional[Hash]) -> Path:
+        """Closure over a URI that downloads a file to a local path and returns the path.
+        The file may be placed anywhere as long as the file will be readable until the
+        program exits.
+
+        If the URI points to a missing file, this MUST raise any Exception that the
+        underlying implementation desires. It MUST NOT return a Path pointing to a
+        non-existent file.
+
+        The Hash may be used to short-circuit a download that would result in downloading
+        a file that does not match the expected hash, but the Downloader need not verify
+        the Hash of the file downloaded after the fact, as that will be performed by
+        default by the Source object.
+        """
+
+
+class DownloadHandler(ty.Protocol):
+    def __call__(self, uri: str) -> ty.Optional[Downloader]:
+        """Returns a Downloader containing the URI if this URI can be handled.  Returns
+        None if this URI cannot be handled.
+        """
+
+
+def _LocalFileHandler(uri: str) -> ty.Optional[Downloader]:
+    if not is_file_uri(uri):
+        return None
+
+    def download_file(hash: ty.Optional[Hash]) -> Path:
+        lpath = path_from_uri(uri)
+        if not lpath.exists():
+            raise FileNotFoundError(lpath)
+        if hash:
+            _check_hash(hash, lpath)
+        return lpath
+
+    return download_file
+
+
+def register_download_handler(key: str, handler: DownloadHandler):
+    # key is not currently used for anything other than avoiding
+    # having duplicates registered for whatever reason.
+    _DOWNLOAD_HANDLERS[key] = handler
+
+
+_DOWNLOAD_HANDLERS: ty.Dict[str, DownloadHandler] = dict()
+register_download_handler("local_file", _LocalFileHandler)
+
+
+def _get_download_handler(uri: str) -> Downloader:
+    for handler in _DOWNLOAD_HANDLERS.values():
+        if downloader := handler(uri):
+            return downloader
+    raise ValueError(f"No SourcePath download handler for uri: {uri}")
+
+
+class SourceHashMismatchError(ValueError):
+    pass
+
+
+def _check_hash(expected_hash: ty.Optional[Hash], path: Path) -> Hash:
+    hash_algo = expected_hash.algo if expected_hash else "sha256"
+    with log.logger_context(hash_for=f"source-{hash_algo}"):
+        computed_hash = filehash(hash_algo, path)
+    if expected_hash and expected_hash != computed_hash:
+        raise SourceHashMismatchError(
+            f"{expected_hash.algo} mismatch for {path};"
+            f" got {computed_hash.bytes!r}, expected {expected_hash.bytes!r}"
+        )
+    return computed_hash
+
+
+@dataclass(frozen=True)
+class Source(os.PathLike):
+    """Source is meant to be a consistent in-memory representation for an abstract,
+    **read-only** source of data that may not be present locally when an application
+    starts.
+
+    A Source uses `os.PathLike` (`__fspath__`) to support transparent `open(src)` calls,
+    so in many cases it will be a drop-in replacement for Path or str filenames. If you
+    need an actual Path object, you can call `path()` to get one, but you should prefer to
+    defer this until the actual location of use.
+
+    By 'wrapping' read-only data in these objects, we can unify the code around how we
+    unwrap and use them, which should allow us to more easily support different execution
+    environments and sources of data.
+
+    For instance, a Source could be a file on disk, but it could also be a file in
+    ADLS.
+
+    Furthermore, libraries which build on top of this one may use this representation to
+    identify opportunities for optimization, by representing the Source in a stable
+    and consistent format that allows different underlying data sources to fulfill the
+    request for the data based on environmental context. A library could choose to
+    transparently transform a local-path-based Source into a Source representing a
+    remote file, without changing the semantics of the object as observed by the code.
+
+    One reason a Hash is part of the interface is so that libraries interacting with the
+    object can use the hash as a canonical 'name' for the data, if one is available.
+
+    Another reason is that we can add a layer of consistency checking to data we're
+    working with, at the cost of a few compute cycles. Since Sources are meant to represent
+    read-only data, the Hash is a meaningful and persistent marker of data identity.
+
+    Do not call its constructor in application code. Use `from_file` or `from_uri` instead.
+    """
+
+    uri: str
+    hash: ty.Optional[Hash] = None
+    # hash and equality are based only on the _identity_ of the object,
+    # not on the other properties that provide some caching functionality.
+
+    @property
+    def cached_path(self) -> ty.Optional[Path]:
+        """This is part of the public interface as far as checking to see whether a file
+        is already present locally, but its existence and value is not part of equality or
+        the hash for this class - it exists purely as an optimization.
+        """
+        return getattr(self, "__cached_path", None)
+
+    def _set_cached_path(self, lpath: ty.Optional[Path]):
+        """protected interface for setting a cached Path since the attribute is not
+        available via the constructor.
+        """
+        super().__setattr__("__cached_path", lpath)  # this works around dataclass.frozen.
+        # https://noklam.github.io/blog/posts/2022-04-22-python-dataclass-partiala-immutable.html
+
+    def path(self) -> Path:
+        """Any Source can be turned into a local file path.
+
+        Remember that the resulting data is meant to be read-only. If you want to mutate
+        the data, you should first make a copy.
+
+        If not already present locally, this will incur a one-time download. Then, if the
+        Source has a Hash, the Hash will be validated against the downloaded file, and a
+        failure will raise SourceHashMismatchError.
+        """
+        if self.cached_path is None or not self.cached_path.exists():
+            lpath = _get_download_handler(self.uri)(self.hash)
+            # path() used to be responsible for checking the hash, but since we pass it to the downloader,
+            # it really makes more sense to allow the downloader to decide how to verify its own download,
+            # and we don't want to duplicate any effort that it may have already put in.
+            self._set_cached_path(lpath)
+
+        assert self.cached_path and self.cached_path.exists()
+        return self.cached_path
+
+    def __fspath__(self) -> str:
+        return os.fspath(self.path())
+
+
+# Creation from local Files or from remote URIs
+
+
+def from_file(filename: StrOrPath, hash: ty.Optional[Hash] = None, uri: str = "") -> Source:
+    """Create a read-only Source from a local file that already exists.
+
+    If URI is passed, the local file will be read and hashed, but the final URI in the
+    Source will be the one provided explicitly. NO UPLOAD IS PERFORMED. It is your
+    responsibility to ensure that your file has been uploaded to the URI you provide.
+    """
+    path = path_from_uri(filename) if isinstance(filename, str) else filename
+    assert isinstance(path, Path)
+    if not path.exists():
+        raise FileNotFoundError(path)
+
+    if uri:
+        src = from_uri(uri, _check_hash(hash, path))
+    else:
+        src = Source(to_uri(path), _check_hash(hash, path))
+    src._set_cached_path(path)  # internally, it's okay to hack around immutability.
+    return src
+
+
+class FromUri(ty.Protocol):
+    def __call__(self, hash: ty.Optional[Hash]) -> Source:
+        """Closure over a URI that creates a Source from a URI.
+
+        The Hash may be used to short-circuit creation that would result in creating
+        a Source that cannot match the expected Hash, but this is not required,
+        and the hash will be included in the Source object regardless, and will
+        be validated (if non-nil) at the time of source data access.
+        """
+
+
+class FromUriHandler(ty.Protocol):
+    def __call__(self, uri: str) -> ty.Optional[FromUri]:
+        """Returns a FromUri object containing the URI if this URI can be handled.  Returns
+        None if this URI cannot be handled.
+        """
+
+
+def register_from_uri_handler(key: str, handler: FromUriHandler):
+    """If a library wants to customize how Sources are created from URIs that it handles,
+    it can register a handler here.
+    """
+    # key is not currently used for anything other than avoiding
+    # having duplicates registered for whatever reason.
+    _FROM_URI_HANDLERS[key] = handler
+
+
+_FROM_URI_HANDLERS: ty.Dict[str, FromUriHandler] = dict()
+register_from_uri_handler(
+    "local_file", lambda uri: partial(from_file, path_from_uri(uri)) if is_file_uri(uri) else None
+)
+
+
+def from_uri(uri: str, hash: ty.Optional[Hash] = None) -> Source:
+    """Create a read-only Source from a URI. The data should already exist at this remote
+    URI, although Source itself can make no guarantee that it always represents real data
+    - only that any data it does represent is read-only.
+
+    It may be advantageous for a URI-handling library to register a more specific
+    implementation of this function, if it is capable of determining a Hash for the blob
+    represented by the URI without downloading the blob.
+    """
+    for handler in _FROM_URI_HANDLERS.values():
+        if from_uri_ := handler(uri):
+            return from_uri_(hash)
+    return Source(uri=uri, hash=hash)

thds/core/source_serde.py

@@ -0,0 +1,104 @@
+# this should later get promoted somewhere, probably
+import json
+import typing as ty
+from functools import partial
+from pathlib import Path
+
+from thds.core import files, hashing, log, source, types
+
+_SHA256_B64 = "sha256b64"
+_MD5_B64 = "md5b64"
+
+logger = log.getLogger(__name__)
+
+
+def _from_sha256b64(d: dict) -> ty.Optional[hashing.Hash]:
+    if "sha256b64" in d:
+        return hashing.Hash(algo="sha256", bytes=hashing.db64(d[_SHA256_B64]))
+    return None
+
+
+def _from_md5b64(d: dict) -> ty.Optional[hashing.Hash]:
+    if "md5b64" in d:
+        return hashing.Hash(algo="md5", bytes=hashing.db64(d[_MD5_B64]))
+    return None
+
+
+HashParser = ty.Callable[[dict], ty.Optional[hashing.Hash]]
+_BASE_PARSERS = (_from_sha256b64, _from_md5b64)
+
+
+def base_parsers() -> ty.Tuple[HashParser, ...]:
+    return _BASE_PARSERS
+
+
+def from_json(
+    json_source: str, hash_parsers: ty.Collection[HashParser] = base_parsers()
+) -> source.Source:
+    d = json.loads(json_source)
+    return source.from_uri(
+        uri=d["uri"],
+        hash=next(filter(None, (p(d) for p in hash_parsers)), None),
+    )
+
+
+def _generic_hash_serializer(
+    algo: str, stringify_hash: ty.Callable[[bytes], str], keyname: str, hash: hashing.Hash
+) -> ty.Optional[dict]:
+    if hash.algo == algo:
+        return {keyname: stringify_hash(hash.bytes)}
+    return None
+
+
+_to_sha256b64 = partial(_generic_hash_serializer, "sha256", hashing.b64, _SHA256_B64)
+_to_md5b64 = partial(_generic_hash_serializer, "md5", hashing.b64, _MD5_B64)
+
+HashSerializer = ty.Callable[[hashing.Hash], ty.Optional[dict]]
+_BASE_HASH_SERIALIZERS: ty.Tuple[HashSerializer, ...] = (_to_md5b64, _to_sha256b64)  # type: ignore
+
+
+def base_hash_serializers() -> ty.Tuple[HashSerializer, ...]:
+    return _BASE_HASH_SERIALIZERS
+
+
+def to_json(
+    source: source.Source, hash_serializers: ty.Collection[HashSerializer] = base_hash_serializers()
+) -> str:
+    hash_dict = (
+        next(filter(None, (ser(source.hash) for ser in hash_serializers if source.hash)), None)
+    ) or dict()
+    return json.dumps(dict(uri=source.uri, **hash_dict))
+
+
+def from_unknown_user_path(path: types.StrOrPath, desired_uri: str) -> source.Source:
+    """Sometimes you may want to load a Source directly from a Path provided by a user.
+
+    It _might_ represent something loadable as a from_json Source, but it might just be a
+    raw file that needs to be loaded with from_file!
+
+    This is a _reasonable_ (but not guaranteed!) way of trying to ascertain which one it
+    is, and specifying where it should live 'remotely' if such a thing becomes
+    necessary.
+
+    Your application might need to implement something more robust if the
+    actual underlying data is likely to be a JSON blob containing the key `uri`, for
+    instance.
+    """
+    with open(path) as readable:
+        try:
+            return from_json(readable.read(4096))
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            return source.from_file(path, uri=desired_uri)
+
+
+def write_to_json_file(source: source.Source, local_file: Path) -> bool:
+    """Write the canonical JSON serialization of the Source to a file."""
+    local_file.parent.mkdir(parents=True, exist_ok=True)
+    previous_source = local_file.read_text() if local_file.exists() else None
+    new_source = to_json(source) + "\n"
+    if new_source != previous_source:
+        with files.atomic_text_writer(local_file) as f:
+            logger.info(f"Writing {source} to {local_file}")
+            f.write(new_source)
+            return True
+    return False

thds/core/sqlite/__init__.py

@@ -0,0 +1,21 @@
+from . import connect, copy, ddl, functions, index, read, sqlmap, upsert  # noqa: F401
+from .merge import merge_databases  # noqa: F401
+from .meta import (  # noqa: F401
+    debug_errors,
+    list_tables,
+    preload_sources,
+    primary_key_cols,
+    table_name_from_path,
+    table_source,
+)
+from .structured import StructTable, struct_table_from_source  # noqa: F401
+from .types import (  # noqa: F401
+    AnyDbTableSrc,
+    DbAndTable,
+    DbAndTableP,
+    TableMaster,
+    TableSource,
+    maybe_t,
+    resolve_lazy_db_and_table,
+)
+from .write import make_mapping_writer, write_mappings  # noqa: F401

thds/core/sqlite/connect.py

@@ -0,0 +1,33 @@
+import contextlib
+import os
+import sqlite3
+import typing as ty
+
+from thds.core import scope
+
+from .functions import register_functions_on_connection
+from .types import Connectable
+
+
+def row_connect(path: ty.Union[str, os.PathLike]) -> sqlite3.Connection:
+    """Get a connection to a row database"""
+    conn = sqlite3.connect(os.fspath(path), isolation_level=None)  # autocommit
+    conn.row_factory = sqlite3.Row
+    return register_functions_on_connection(conn)
+
+
+autoconn_scope = scope.Scope("sqlite3.autoconn")
+
+
+def autoconnect(connectable: Connectable) -> sqlite3.Connection:
+    """Will automatically commit when it hits the autoconn_scope.bound, but only if
+    the connectable was not already a connection.
+    """
+    if isinstance(connectable, sqlite3.Connection):
+        return connectable
+
+    return autoconn_scope.enter(
+        contextlib.closing(  # close the connection when we exit the scope
+            row_connect(os.fspath(connectable))
+        )
+    )

thds/core/sqlite/copy.py

@@ -0,0 +1,35 @@
+"""Utility for copying a table from one connection to another."""
+
+from ..log import getLogger
+from .connect import autoconn_scope, autoconnect
+from .types import Connectable
+
+logger = getLogger(__name__)
+
+
+@autoconn_scope.bound
+def table(source: Connectable, table_name: str, dest: Connectable) -> None:
+    """Copy a table from one connection to another, including its table definition.
+
+    If you can do this using ATTACH instead, that will be faster because it involves
+    no Python code running in the loop.
+    """
+    source_conn = autoconnect(source)
+    dest_conn = autoconnect(dest)
+
+    source_table_sql = source_conn.execute(
+        f"SELECT sql FROM sqlite_master WHERE name = '{table_name}'"
+    ).fetchone()[0]
+
+    dest_conn.execute(source_table_sql)
+
+    src_data = source_conn.execute(f"SELECT * FROM {table_name}")
+
+    dest_conn.execute("BEGIN TRANSACTION;")
+    while True:
+        data = src_data.fetchmany(1000)
+        if not data:
+            break
+        placeholders = ", ".join(["?"] * len(data[0]))
+        dest_conn.executemany(f"INSERT INTO {table_name} VALUES ({placeholders});", data)
+    dest_conn.execute("COMMIT;")

thds/core/sqlite/ddl.py

@@ -0,0 +1,4 @@
+def drop(full_name: str, is_view: bool = False) -> str:
+    """Drop a table or view."""
+    table_or_view = "TABLE" if not is_view else "VIEW"
+    return f"DROP {table_or_view} IF EXISTS {full_name};"