splime 0.1.2__py3-none-any.whl

spl/core/entities/artifact.py

@@ -0,0 +1,141 @@
+import ast
+import hashlib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Generator
+
+import yaml
+
+from spl.core.ir.common import DBase
+from spl.core.ir.parse import _branch, ir_parse
+from spl.core.ir.unparse import ir_unparse
+
+
+_HASH_CHUNK_SIZE = 1024 * 1024
+_SHA256_HEX_DIGITS = set('0123456789abcdefABCDEF')
+
+
+def _validate_non_empty_string(name: str, value: str) -> None:
+    if not isinstance(value, str):
+        raise TypeError('artifact ref {} must be a string'.format(name))
+    if not value:
+        raise ValueError('artifact ref {} must be a non-empty string'.format(name))
+
+
+def _validate_sha256(value: str) -> None:
+    _validate_non_empty_string('sha256', value)
+    if len(value) != 64 or any(c not in _SHA256_HEX_DIGITS for c in value):
+        raise ValueError('artifact ref sha256 must be a 64-character hex string')
+
+
+def _validate_size(value: int) -> None:
+    if type(value) is not int:
+        raise TypeError('artifact ref size must be an integer')
+    if value < 0:
+        raise ValueError('artifact ref size must be non-negative')
+
+
+def _validate_artifact_ref(key: str, uri: str, sha256: str, size: int) -> None:
+    _validate_non_empty_string('key', key)
+    _validate_non_empty_string('uri', uri)
+    _validate_sha256(sha256)
+    _validate_size(size)
+
+
+@dataclass(frozen = True)
+class ArtifactRef:
+    """Reference to a materialized artifact on the local filesystem."""
+
+    key: str
+    uri: str
+    sha256: str
+    size: int
+
+    def __post_init__(self) -> None:
+        _validate_artifact_ref(
+            key = self.key,
+            uri = self.uri,
+            sha256 = self.sha256,
+            size = self.size)
+
+
+@dataclass(frozen = True)
+class DArtifactRef(DBase):
+    """Serialized ArtifactRef value for pipeline YAML."""
+
+    key: str
+    uri: str
+    sha256: str
+    size: int
+
+    def __post_init__(self) -> None:
+        _validate_artifact_ref(
+            key = self.key,
+            uri = self.uri,
+            sha256 = self.sha256,
+            size = self.size)
+
+
+yaml.add_representer(
+    DArtifactRef,
+    lambda dumper, data: dumper.represent_mapping('!DArtifactRef', data.__dict__))
+
+
+def _construct_dartifact_ref(loader: Any, node: Any) -> DArtifactRef:
+    return DArtifactRef(**loader.construct_mapping(node))
+
+
+yaml.add_constructor(
+    '!DArtifactRef',
+    _construct_dartifact_ref)
+
+
+def compute_sha256(path: Path) -> str:
+    """Compute a file's SHA-256 digest with chunked reads."""
+
+    digest = hashlib.sha256()
+    with path.open('rb') as f:
+        for chunk in iter(lambda: f.read(_HASH_CHUNK_SIZE), b''):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _mk_empty_dependencies(frame_offset: int) -> Generator[DBase]:
+    yield from ()
+
+
+@ir_parse.register(
+    lambda x: isinstance(x, ArtifactRef))
+def _ir_parse__artifact_ref(
+        x: ArtifactRef,
+        name: str | None = None) -> _branch:
+    return _branch(
+        x,
+        lambda: DArtifactRef(
+            key = x.key,
+            uri = x.uri,
+            sha256 = x.sha256,
+            size = x.size),
+        _mk_empty_dependencies)
+
+
+@ir_unparse.register(
+    lambda x: isinstance(x, DArtifactRef))
+def _ir_unparse__artifact_ref(x: DArtifactRef, source: Path) -> Generator[ast.stmt]:
+    yield ast.Assign(
+        targets = [ast.Name(id = '_link_to', ctx = ast.Store())],
+        value = ast.Call(
+            func = ast.Name(id = 'ArtifactRef', ctx = ast.Load()),
+            keywords = [
+                ast.keyword(
+                    arg = 'key',
+                    value = ast.Constant(value = x.key)),
+                ast.keyword(
+                    arg = 'uri',
+                    value = ast.Constant(value = x.uri)),
+                ast.keyword(
+                    arg = 'sha256',
+                    value = ast.Constant(value = x.sha256)),
+                ast.keyword(
+                    arg = 'size',
+                    value = ast.Constant(value = x.size))]))

spl/core/entities/control.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+import yaml
+
+from spl.core.ir.common import DBase
+from spl.core.ir.unparse import ir_unparse
+
+
+@dataclass(frozen = True)
+class DSPLSelfImport(DBase):
+    name: str
+
+yaml.add_representer(
+    DSPLSelfImport,
+    lambda dumper, data: dumper.represent_mapping('!DSPLSelfImport', data.__dict__))
+
+yaml.add_constructor(
+    '!DSPLSelfImport',
+    lambda loader, node: DSPLSelfImport(**loader.construct_mapping(node)))
+
+
+@dataclass(frozen = True)
+class DSPLImport(DBase):
+    path: str
+    name: str
+
+yaml.add_representer(
+    DSPLImport,
+    lambda dumper, data: dumper.represent_mapping('!DSPLImport', data.__dict__))
+
+yaml.add_constructor(
+    '!DSPLImport',
+    lambda loader, node: DSPLImport(**loader.construct_mapping(node)))
+
+
+@ir_unparse.register(lambda x: isinstance(x, DSPLSelfImport))
+def _ir_unparse__spl_self_import(x: DSPLSelfImport, source: Path):
+    yield from []
+
+
+@ir_unparse.register(lambda x: isinstance(x, DSPLImport))
+def _ir_unparse__spl_import(x: DSPLImport, source: Path):
+    yield from []
+

spl/core/entities/distribution.py

@@ -0,0 +1,65 @@
+import importlib
+import logging
+from dataclasses import dataclass
+from importlib.metadata import PackageNotFoundError, packages_distributions
+from itertools import chain
+from pathlib import Path
+from types import ModuleType
+from typing import Generator
+
+import yaml
+
+from spl.core.ir.common import DBase
+from spl.core.ir.unparse import ir_unparse
+
+
+@dataclass(frozen = True)
+class DDistribution(DBase):
+    package: str
+    version: str
+
+    def __lt__(self, other):
+        return (self.package, self.version) < (other.package, other.version)
+
+yaml.add_representer(
+    DDistribution,
+    lambda dumper, data: dumper.represent_mapping('!DDistribution', data.__dict__))
+
+yaml.add_constructor(
+    '!DDistribution',
+    lambda loader, node: DDistribution(**loader.construct_mapping(node)))
+
+
+def get_dependencies_from_distribution(module: ModuleType) -> Generator[DDistribution]:
+    distributions = packages_distributions()
+    if (package := module.__package__):
+        for x in set(distributions[package.split('.')[0]]):
+            yield DDistribution(
+                package = x,
+                version = importlib.metadata.version(x))
+
+
+def validate_distributions(deps, source) -> None:
+    distributions = sorted(set(filter(
+        lambda x: isinstance(x, DDistribution),
+        chain.from_iterable([dependencies for (_, dependencies) in deps]))))
+
+    for x in distributions:
+        try:
+          if (version := importlib.metadata.version(x.package)) != x.version:
+                logging.warning(
+                    '{}: distribution mismatch: {} == {} (actual {})'.format(
+                        source,
+                        x.package,
+                        x.version, version))
+        except PackageNotFoundError:
+            logging.warning(
+                '{}: distribution is not found: {} == {}'.format(
+                    source,
+                    x.package,
+                    x.version))
+
+
+@ir_unparse.register(lambda x: isinstance(x, DDistribution))
+def _ir_unparse__distribution(x: DDistribution, source: Path):
+    yield from []

spl/core/entities/function.py

@@ -0,0 +1,254 @@
+import ast
+import builtins
+import dis
+import inspect
+import sys
+import typing
+from dataclasses import dataclass
+from itertools import chain
+from pathlib import Path
+from types import CodeType, FunctionType
+from typing import Annotated, Generator
+
+import yaml
+
+import spl.core.entities.node as m_node
+from spl.core.entities.control import DSPLImport
+from spl.core.entities.node import DEFAULT_PORT, InputPort, OutputPort
+from spl.core.ir.common import DBase
+from spl.core.ir.parse import _attach, _branch, ir_parse
+from spl.core.ir.unparse import ir_unparse
+
+LOCATION_DUNDER_NAME = '__spl_location__'
+METADATA_DUNDER_NAME = '__spl_metadata__'
+
+
+@dataclass
+class _body:  # noqa: N801
+    value: str
+
+yaml.add_representer(
+    _body,
+    lambda dumper, data: dumper.represent_scalar('tag:yaml.org,2002:str', data.value, style = '|'))
+
+
+@dataclass(frozen = True)
+class DFunction(DBase):
+    name: str
+    body: str
+    inputs: list[InputPort]
+    outputs: list[OutputPort] | None
+
+    def __hash__(self):
+        return hash((self.name, self.body, tuple(self.inputs), tuple(self.outputs)))
+
+    def __repr__(self):
+        return '{}({})'.format(self.__class__.__name__, repr(self.name))
+
+
+yaml.add_representer(
+    DFunction,
+    lambda dumper, data: dumper.represent_mapping('!DFunction', {
+        'name': data.name,
+        'inputs': list(map(
+            lambda x: {
+                'name': x.name,
+                'type': x.typ_,
+                'default': x.default}, data.inputs)),
+        'outputs': None if data.outputs is None else list(map(
+            lambda x: {
+                'name': x.name,
+                'type': x.typ_}, data.outputs)),
+        'body': _body(data.body)}))
+
+
+yaml.add_constructor(
+    '!DFunction',
+    lambda loader, node: (lambda data: DFunction(
+        name = data['name'],
+        body = data['body'],
+        inputs = [
+            InputPort(
+                name = x['name'],
+                typ_ = x['type'],
+                default = x.get('default'))
+            for x in data['inputs']],
+        outputs = [
+            OutputPort(
+                name = x['name'],
+                typ_ = x['type'])
+            for x in data['outputs']]))
+    (loader.construct_mapping(node, deep = True)))
+
+
+def get_dependency_names_from_bytecode(f: FunctionType | CodeType):
+    for x in dis.Bytecode(f):
+        match x.opname:
+            case 'LOAD_GLOBAL':
+                yield x.argval
+            case 'LOAD_NAME':
+                yield x.argval
+            case 'LOAD_CONST':
+                if isinstance(x.argval, CodeType):
+                    yield from get_dependency_names_from_bytecode(x.argval)
+
+def get_dependencies_from_bytecode(
+        frame_offset,
+        f: FunctionType | CodeType):
+    g = sys._getframe(1 + frame_offset).f_globals
+
+    names = sorted(filter(
+        lambda x: x not in vars(builtins),
+        set(get_dependency_names_from_bytecode(f))))
+
+    if missing_names := (set(names) - set(g.keys())):
+        raise ValueError('missing names: {}'.format(', '.join(sorted(missing_names))))
+
+    for name in names:
+        yield ir_parse(g[name], name)
+
+
+def get_dependencies_from_ast(frame_offset, tree: ast.FunctionDef):
+    for value in map(ast.unparse, filter(None, [x.annotation for x in tree.args.args])):
+        yield from get_dependencies_from_bytecode(
+            frame_offset + 1,
+            value)
+    for value in map(ast.unparse, tree.args.defaults):
+        yield from get_dependencies_from_bytecode(
+            frame_offset + 1,
+            value)
+
+
+def serialize_function_output(func: FunctionType, tree: ast.FunctionDef):
+    name = DEFAULT_PORT
+
+    return_type = func.__annotations__.get('return')
+    if (return_type is not None) and (typing.get_origin(return_type) == Annotated):
+        (_, name, *_) = typing.get_args(return_type)
+
+    return [
+        OutputPort(
+            name = name,
+            typ_ = ast.unparse(tree.returns) if (tree.returns is not None) else None)]
+
+
+def serialize_function(func: FunctionType, tree: ast.FunctionDef | None = None):
+    if tree is None:
+        [tree] = ast.parse(inspect.getsource(func)).body
+
+    args = tree.args
+
+    return DFunction(
+        name = tree.name,
+        body = ast.unparse(tree.body),
+        inputs = [
+            InputPort(
+                name = a.arg,
+                typ_ = ast.unparse(a.annotation) if a.annotation is not None else None,
+                default = ast.unparse(d) if d is not None else None)
+            for a, d in zip(
+                    args.args,
+                    [*[None] * (len(args.args) - len(args.defaults)), *args.defaults],
+                    strict = True)],
+        outputs = serialize_function_output(func, tree))
+
+
+@ir_parse.register(
+    lambda x: (isinstance(x, FunctionType) and (getattr(x, '__module__', None) == '__main__')))
+def _ir_parse__function(
+        x: FunctionType,
+        name: str | None = None):
+
+    if hasattr(x, LOCATION_DUNDER_NAME):
+        # We imported this function using SPL, using it's metadata.
+        return _attach(chain([DSPLImport(*getattr(x, LOCATION_DUNDER_NAME))]))
+
+    [tree] = ast.parse(inspect.getsource(x)).body
+    return _branch(
+        x,
+        lambda: serialize_function(x, tree),
+        lambda frame_offset: chain(
+            get_dependencies_from_bytecode(frame_offset, x),
+            get_dependencies_from_ast(frame_offset, tree)))
+
+
+def get_function_metadata(func: FunctionType):
+    if hasattr(func, METADATA_DUNDER_NAME):
+        # We imported this function using SPL, using it's metadata.
+        return getattr(func, METADATA_DUNDER_NAME)
+    return serialize_function(func)
+
+
+@ir_unparse.register(lambda x: isinstance(x, DFunction))
+def _ir_unparse__function(x: DFunction, source: Path) -> Generator[ast.stmt]:
+    yield ast.FunctionDef(
+        name = x.name,
+        body = ast.parse(x.body).body,
+        # TODO: add support for multiple outputs
+        returns = ast.parse(x.outputs[0].typ_, mode = 'eval').body if x.outputs[0].typ_ is not None else None,
+        args = ast.arguments(
+            args = [
+                ast.arg(
+                    arg = port.name,
+                    annotation = ast.parse(port.typ_, mode = 'eval').body if port.typ_ is not None else None)
+                for port in x.inputs],
+            defaults = [
+                ast.parse(port.default, mode = 'eval').body
+                for port in x.inputs
+                if port.default is not None]))
+
+    # Importing helpers
+    yield ast.ImportFrom(
+        module = m_node.__name__,
+        names = [
+            ast.alias(name = 'InputPort'),
+            ast.alias(name = 'OutputPort')])
+
+    yield ast.ImportFrom(
+        module = __name__,
+        names = [
+            ast.alias(name = 'DFunction')])
+
+    # Marking function as spl-imported
+    yield ast.Expr(value = ast.Call(
+        func = ast.Name(id = 'setattr', ctx = ast.Load()),
+        args = [
+            ast.Name(id = x.name, ctx = ast.Load()),
+            ast.Constant(value = LOCATION_DUNDER_NAME),
+            ast.Tuple(elts = [
+                ast.Constant(value = str(source.absolute())),
+                ast.Constant(value = x.name)])]))
+
+    # Adding metadata
+    yield ast.Expr(value = ast.Call(
+        func = ast.Name(id = 'setattr', ctx = ast.Load()),
+        args = [
+            ast.Name(id = x.name, ctx = ast.Load()),
+            ast.Constant(value = METADATA_DUNDER_NAME),
+            ast.Call(
+                func = ast.Name(id = 'DFunction', ctx = ast.Load()),
+                keywords = [
+                    ast.keyword(arg = 'name', value = ast.Constant(value = x.name)),
+                    ast.keyword(arg = 'body', value = ast.Constant(value = x.body)),
+                    ast.keyword(arg = 'inputs', value = ast.List(
+                        elts = [
+                            ast.Call(
+                                func = ast.Name(id = 'InputPort', ctx = ast.Load()),
+                                keywords = [
+                                    ast.keyword(arg = 'name', value = ast.Constant(value = port.name)),
+                                    ast.keyword(arg = 'typ_', value = ast.Constant(value = port.typ_)),
+                                    ast.keyword(
+                                        arg = 'default',
+                                        value = ast.Constant(value = port.default))])
+                            for port in x.inputs],
+                        ctx = ast.Load())),
+
+                    ast.keyword(arg = 'outputs', value = ast.List(
+                        elts = [
+                            ast.Call(
+                                func = ast.Name(id = 'OutputPort', ctx = ast.Load()),
+                                keywords = [
+                                    ast.keyword(arg = 'name', value = ast.Constant(value = port.name)),
+                                    ast.keyword(arg = 'typ_', value = ast.Constant(value = port.typ_))])
+                            for port in x.outputs],
+                        ctx = ast.Load()))])]))