app/backend/fluksio/sdk/imports.py

"""What code a node actually reaches, so a cache key can name it.

A node's stored body is a shim that imports the real function, so it says
nothing about what that function calls into. Hashing the whole repository
instead says too much: it moves when a notebook two directories away does,
which makes the digest neither necessary nor sufficient for "this run's
numbers came from that code".

This walks the imports out from the module a node's function lives in and
stops at anything installed or standard, so what is left is the project's own
code — including a package installed editable from its own checkout, whose
files are its source rather than a copy under `site-packages`.

The walk is static: it reads the imports out of each file rather than running
it. A module imported by a name the code computes is therefore not followed,
which is the same limit the yield check has, and for the same reason.
"""

from __future__ import annotations

import ast
import hashlib
import importlib.util
import sys
import sysconfig
from collections.abc import Iterator
from pathlib import Path
from typing import Any

__all__ = ["digest_of", "reached"]


#: Never the project's own code, however it happens to be installed. Every
#: node's module imports the authoring API, and from a checkout or an editable
#: install that is a source tree like any other — so without this, a node's
#: digest is mostly this tool, and upgrading it invalidates every cache entry
#: anybody has.
_SELF = frozenset({"fluksio", "fluksio_worker"})


def _foreign() -> tuple[str, ...]:
    """Where code that is not the project's lives.

    The standard library and whatever `pip install` copied. An editable
    install is deliberately not here: its files are the checkout, and a
    checkout is somebody's code.
    """
    roots = {
        path
        for key in ("stdlib", "platstdlib", "purelib", "platlib")
        if (path := sysconfig.get_paths().get(key))
    }
    roots.update(
        entry
        for entry in sys.path
        if entry.endswith(("site-packages", "dist-packages"))
    )
    return tuple(sorted(roots))


def _project_file(name: str, foreign: tuple[str, ...]) -> Path | None:
    """The file a dotted module is, if it is the project's own python."""
    module = sys.modules.get(name)
    origin = getattr(module, "__file__", None) if module is not None else None
    if origin is None:
        try:
            spec = importlib.util.find_spec(name)
        except (ImportError, ValueError, AttributeError):
            return None
        origin = spec.origin if spec is not None else None
    if not origin or not origin.endswith(".py"):
        return None
    return None if origin.startswith(foreign) else Path(origin)


def _package_of(name: str, path: Path) -> str:
    """The package a relative import inside this module resolves against."""
    return name if path.name == "__init__.py" else name.rpartition(".")[0]


def _imported(tree: ast.AST, package: str) -> Iterator[str]:
    """Every module name a file imports, relative ones resolved.

    `from pkg import thing` yields both `pkg` and `pkg.thing`, since only
    trying to resolve them says which of the two `thing` is.
    """
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                yield alias.name
        elif isinstance(node, ast.ImportFrom):
            base = node.module or ""
            if node.level:
                parts = package.split(".") if package else []
                parts = parts[: len(parts) - (node.level - 1)]
                base = ".".join([*parts, base]) if base else ".".join(parts)
            if not base:
                continue
            yield base
            for alias in node.names:
                yield f"{base}.{alias.name}"


def reached(fn: Any) -> dict[str, str]:
    """The project modules ``fn``'s own module reaches, itself included.

    Dotted name to the file it is on this machine — the name is what the
    digest is keyed on, so an engine hashing the same code arrives at the same
    answer; the path is how it finds the file.
    """
    start = getattr(fn, "__module__", "")
    if not start:
        return {}
    foreign = _foreign()
    found: dict[str, str] = {}
    queue = [start]
    while queue:
        name = queue.pop()
        if name in found or name.partition(".")[0] in _SELF:
            continue
        path = _project_file(name, foreign)
        if path is None:
            continue
        try:
            tree = ast.parse(path.read_bytes())
        except (OSError, SyntaxError, ValueError):
            continue
        found[name] = str(path)
        queue.extend(_imported(tree, _package_of(name, path)))
        # The packages above it, whose `__init__` runs on the way to importing
        # this module and is as much a part of what it does as an import is.
        parts = name.split(".")[:-1]
        queue.extend(".".join(parts[: index + 1]) for index in range(len(parts)))
    return found


def digest_of(files: dict[str, str]) -> str:
    """What those modules hash to now, or empty if one is not readable here.

    All or nothing: a digest over the half of the code an engine happens to
    have would be a confident answer to a question it cannot see.
    """
    if not files:
        return ""
    digest = hashlib.sha256()
    for name, path in sorted(files.items()):
        try:
            body = Path(path).read_bytes()
        except OSError:
            return ""
        digest.update(name.encode())
        digest.update(b"\0")
        digest.update(body)
        digest.update(b"\0")
    return digest.hexdigest()