procfunc 0.30.0__py3-none-any.whl

procfunc/transpiler/main.py

@@ -0,0 +1,299 @@
+import argparse
+import itertools
+import logging
+from pathlib import Path
+from typing import Callable, Literal
+
+import bpy
+import pandas as pd
+
+import procfunc as pf
+from procfunc import compute_graph as cg
+from procfunc import transforms as tr
+from procfunc.nodes import NODE_OPERATOR_TABLE
+from procfunc.util import pytree
+from procfunc.util.teardown import skip_teardown_on_exit
+
+from .bpy_to_computegraph import (
+    ParseMemo,
+    parse_material,
+    parse_node_tree,
+    parse_object,
+)
+from .codegen import default_func_resolution_map, to_python
+
+logging.basicConfig(
+    format="[%(asctime)s.%(msecs)03d] [%(module)s] [%(levelname)s] | %(message)s",
+    datefmt="%H:%M:%S",
+    level=logging.INFO,
+)
+
+logger = logging.getLogger(__name__)
+
+MODIFIERS_MANIFEST: pd.DataFrame | None = None
+NODES_MANIFEST: pd.DataFrame | None = None
+
+
+BPY_COLLECTIONS = [
+    "scene",
+    "object",
+    "material",
+    # "collections",
+]
+
+RETURN_TYPES = {
+    "object": pf.MeshObject,
+    "material": pf.Material,
+    "collection": pf.Collection,
+}
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("input", type=Path)
+
+    parser.add_argument("--materials", type=str, default=[], nargs="+")
+    parser.add_argument("--objects", type=str, default=[], nargs="+")
+    parser.add_argument("--node_trees", type=str, default=[], nargs="+")
+
+    parser.add_argument("--output", type=Path, default=None)
+    parser.add_argument("--transforms", type=str, default=[], nargs="+")
+
+    parser.add_argument(
+        "--object_mode",
+        type=str,
+        choices=["monkey", "active", "named"],
+        default="monkey",
+    )
+    parser.add_argument("--no_version_comments", action="store_true")
+
+    parser.add_argument(
+        "-d",
+        "--debug",
+        action="store_const",
+        dest="loglevel",
+        const=logging.DEBUG,
+        default=logging.INFO,
+    )
+
+    parser.add_argument("--add_line_comments", action="store_true")
+    parser.add_argument("--include_object_materials", action="store_true")
+
+    return parser
+
+
+def parse_target(
+    target: bpy.types.Object | bpy.types.Material | bpy.types.NodeTree,
+    memo: ParseMemo,
+    object_mode: Literal["monkey", "active", "named"] = "monkey",
+    include_set_material: bool = True,
+) -> cg.ComputeGraph:
+    match target:
+        # case pf.Scene:
+        #    return parse_scene(bpy.context.scene, memo)
+        case bpy.types.Object():
+            return parse_object(
+                target,
+                memo,
+                object_mode=object_mode,
+                include_set_material=include_set_material,
+            )
+        case bpy.types.Material():
+            return parse_material(target, memo)
+        case bpy.types.NodeTree():
+            graph, _ = parse_node_tree(target, memo)
+            return graph
+        case x:
+            raise ValueError(f"Invalid {x=} {type(x)=}")
+
+
+def _find_target_str(target: str, col: bpy.types.bpy_prop_collection) -> list:
+    if target.isdigit():
+        return [col[int(target)]]
+    elif "*" in target:
+        assert target.count("*") == 1
+        assert target.endswith("*") == 0
+        prefix = target.replace("*", "")
+        return [v for k, v in col.items() if k.startswith(prefix)]
+    else:
+        matches = [v for k, v in col.items() if k == target]
+        if len(matches) != 1:
+            available = " ".join(col.keys())
+            raise ValueError(
+                f"Found {len(matches)} targets for {target=} in {available=}"
+            )
+        return [matches[0]]
+
+
+def transpile_targets(
+    targets: list[bpy.types.Object | bpy.types.Material | bpy.types.NodeTree],
+    transforms: list[Callable[[list[cg.ComputeGraph]], list[cg.ComputeGraph]]],
+    object_mode: Literal["monkey", "active", "named"] = "monkey",
+    add_version_comment: bool = True,
+    add_line_comments: bool = False,
+    include_set_material: bool = True,
+) -> str:
+    memo = ParseMemo()
+
+    result_graphs = [
+        parse_target(
+            target,
+            memo,
+            object_mode=object_mode,
+            include_set_material=include_set_material,
+        )
+        for target in targets
+    ]
+
+    for tfunc in transforms:
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug(f"Applying transform {tfunc} to {result_graphs=}")
+        result_graphs = tfunc(result_graphs)
+
+    vec = cg.FunctionCallNode(pf.nodes.shader.coord, (), {})
+    vec = cg.GetAttributeNode(vec, "generated")
+
+    result_calls = []
+    for result_graph in result_graphs:
+        kwargs = {}
+        if result_graph.name.startswith("material_"):
+            kwargs["vector"] = vec
+
+        return_type = RETURN_TYPES.get(type(result_graph), None)
+        call_node = cg.SubgraphCallNode(
+            subgraph=result_graph,
+            args=(),
+            kwargs=kwargs,
+            metadata={"known_value_type": return_type},
+        )
+        result_calls.append(call_node)
+
+    toplevel_name = "toplevel_invoke"
+    graph = cg.ComputeGraph(
+        inputs=pytree.PyTree(()),
+        outputs=pytree.PyTree({"result": result_calls}),
+        name=toplevel_name,
+        metadata={},
+    )
+    assert isinstance(graph, cg.ComputeGraph), graph
+
+    func_resolution, import_lines = default_func_resolution_map(graph)
+    for oprow in NODE_OPERATOR_TABLE:
+        func_resolution[oprow.pf_func] = oprow.operator_type
+
+    python = to_python(
+        graph,
+        func_resolution=func_resolution,
+        import_lines=import_lines,
+        add_version_comment=add_version_comment,
+        add_line_comments=add_line_comments,
+    )
+
+    return python
+
+
+def inject_into_python_file(python: list[str], file: Path, target_name: str):
+    """
+    If def <target_name>(): is found in the python file, replace it with the python code.
+    If not, add it to end of the file.
+    """
+
+    text = file.read_text()
+    lines = text.splitlines()
+
+    start_line = next(
+        (i for i, line in enumerate(lines) if line.startswith(f"def {target_name}")),
+        None,
+    )
+
+    if start_line is None:
+        lines.append
+
+
+def _targets_from_args(
+    args: argparse.Namespace,
+) -> list[bpy.types.Object | bpy.types.Material | bpy.types.NodeTree]:
+    targets = []
+    targets += [
+        _find_target_str(target, bpy.data.materials) for target in args.materials
+    ]
+    for target in args.objects:
+        if target == "ACTIVE":
+            if bpy.context.active_object is None:
+                raise ValueError(
+                    "--objects ACTIVE specified but no active object in scene"
+                )
+            targets.append([bpy.context.active_object])
+        else:
+            targets.append(_find_target_str(target, bpy.data.objects))
+    targets += [
+        _find_target_str(target, bpy.data.node_groups) for target in args.node_trees
+    ]
+
+    return list(itertools.chain.from_iterable(targets))
+
+
+_transforms_map = {
+    "colors_to_hsv_definition": tr.map_graph_list(tr.colors_to_hsv_definition),
+    "infer_distribution_hypercube": (
+        lambda graphs: graphs + [tr.infer_distribution_hypercube(graphs)]
+    ),
+    "infer_nodegroup_distributions": (
+        lambda graphs: tr.infer_nodegroup_distributions(graphs) + graphs
+    ),
+    "extract_materials": tr.extract_materials_from_graphs,
+    # "infer_distribution_polytope": lambda graphs: graphs + [tr.infer_distribution_polytope(graphs)],
+}
+
+
+def main():
+    parser = get_parser()
+    args = parser.parse_args()
+
+    for name in logging.root.manager.loggerDict:
+        logging.getLogger(name).setLevel(args.loglevel)
+
+    pf.ops.file.load_blend(args.input)
+
+    targets = _targets_from_args(args)
+
+    if len(targets) == 0:
+        raise ValueError(
+            f"No targets found for {args.input} {args.materials=} {args.objects=} {args.node_trees=}"
+        )
+
+    logger.info(f"Found targets {[x.name for x in targets]}")
+
+    transforms = [_transforms_map[x] for x in args.transforms]
+    python = transpile_targets(
+        targets,
+        transforms,
+        object_mode=args.object_mode,
+        add_version_comment=not args.no_version_comments,
+        add_line_comments=args.add_line_comments,
+        include_set_material=args.include_object_materials,
+    )
+
+    match args.output:
+        case "print":
+            print(python)
+        case None:
+            pass
+        case x if Path(x).suffix == ".py":
+            x = Path(x)
+            x.parent.mkdir(parents=True, exist_ok=True)
+            print(f"Writing transpiled code to path {x}")
+            with x.open("w") as f:
+                f.write(python)
+        case x if ":" in x:
+            raise NotImplementedError("Not implemented")
+            module_path, target_name = x.split(":")
+            inject_into_python_file(python.splitlines(), Path(module_path), target_name)
+        case _:
+            raise ValueError(f"Invalid output type: {args.output}")
+
+
+if __name__ == "__main__":
+    with skip_teardown_on_exit():
+        main()

procfunc/types.py

@@ -0,0 +1,380 @@
+import logging
+from collections import defaultdict
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Callable, Generic, NamedTuple, TypeAlias, TypeVar
+
+import bpy
+import mathutils
+
+from procfunc.util.pytree import register_pytree_container
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "Asset",
+    "BlenderAsset",
+    "ObjectType",
+    "Object",
+    "CameraObject",
+    "MeshObject",
+    "CurveObject",
+    "EmptyObject",
+    "ArmatureObject",
+    "HairObject",
+    "LatticeObject",
+    "LightObject",
+    "LightProbeObject",
+    "MetaObject",
+    "Material",
+    "Texture",
+    "Image",
+    "Collection",
+    "VolumeObject",
+    "PointCloudObject",
+    "World",
+    "ValueRange",
+]
+
+INVALIDATED_USAGE_ID = -100
+
+
+@dataclass
+class AssetUsageTable:
+    counts: dict[int, int] = field(default_factory=lambda: defaultdict(int))
+
+
+_global_usage_table = AssetUsageTable()
+
+
+def _bpy_data_col_for_asset(item: T) -> bpy.types.bpy_prop_collection:
+    match item.bl_rna.name:
+        case "Object":
+            return bpy.data.objects
+        case "Material":
+            return bpy.data.materials
+        case "Image Texture":
+            return bpy.data.textures
+        case _:
+            raise TypeError(
+                f"{BlenderAsset.__name__} doesnt yet support: {item} with type {item.bl_rna.name}"
+            )
+
+
+class Asset(Generic[T]):
+    def item(self) -> T:
+        raise NotImplementedError("Subclasses of Asset must implement item()")
+
+
+class BlenderAsset(Asset, Generic[T]):
+    """
+    A pythonic wrapper around a blender object,material,texture, or other bpy.data.stuffgoeshere member.
+
+    We reference count the underlying asset and delete it when no python references remain.
+    """
+
+    def __init__(self, item: T):
+        if isinstance(item, Asset):
+            raise ValueError(
+                f"Attempted to wrap {item=} in a second {self.__class__.__name__} wrapper. This is almost certainly not intended"
+            )
+
+        self._item = item
+        _global_usage_table.counts[id(item)] += 1
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug(
+                f"Asset ref {id(self)} created for {self._item.name}. remaining refs: {_global_usage_table.counts[id(self._item)]}"
+            )
+
+        # self._dependencies: list[Asset] = []
+
+    def item(self) -> T:
+        if _global_usage_table.counts[id(self._item)] == INVALIDATED_USAGE_ID:
+            raise ValueError(
+                "Cant access .item() since asset has been explicitly invalidated by a previous operation"
+            )
+
+        return self._item
+
+
+'''
+    def add_dependency(self, dependency: "Asset"):
+        self._dependencies.append(dependency)
+
+    def extend_dependencies(self, other: "Asset"):
+        self._dependencies.extend(other._dependencies)
+
+    def update(self, other: "Asset"):
+        """
+        Updates this asset to point to the same underlying item as other.
+        """
+
+        self._item = other._item
+        self._dependencies = other._dependencies
+        _global_usage_table.counts[id(self._item)] += 1
+
+    def invalidate(self):
+        self._item = None
+        self._dependencies = []
+        _global_usage_table.counts[id(self._item)] = INVALIDATED_USAGE_ID
+
+    def __del__(self):
+        # return
+
+        if self._item is None:
+            return
+        return  # TODO
+        item_id = id(self._item)
+        count = _global_usage_table.counts.get(item_id, 0)
+
+        count -= 1
+
+        # new bad ref counter
+        if count < 0:
+            raise ValueError(
+                f"Asset ref {id(self)} reached negative ref count for item {self._item.name}"
+            )
+
+        _global_usage_table.counts[item_id] = count
+
+        if count == 0:
+            try:
+                bpy_col = _bpy_data_col_for_asset(self._item)
+                if self._item.name in bpy_col:
+                    bpy_col.remove(self._item, do_unlink=True)
+            except ReferenceError:
+                logger.warning(
+                    f"{self.__class__.__name__} __del__ failed - item already deleted"
+                )
+
+            del _global_usage_table.counts[item_id]
+
+        self._item = None
+'''
+
+
+class ObjectType(Enum):
+    ARMATURE = "ARMATURE"
+    CAMERA = "CAMERA"
+    CURVE = "CURVE"
+    EMPTY = "EMPTY"
+    FONT = "FONT"
+    HAIR = "HAIR"
+    LATTICE = "LATTICE"
+    LIGHT = "LIGHT"
+    LIGHT_PROBE = "LIGHT_PROBE"
+    MESH = "MESH"
+    META = "META"
+    POINTCLOUD = "POINTCLOUD"
+    SURFACE = "SURFACE"
+    VOLUME = "VOLUME"
+    # GPENCIL = "GPENCIL"
+    # GREASEPENCIL = "GREASEPENCIL"
+    # SPEAKER = "SPEAKER"
+
+
+Vector: TypeAlias = mathutils.Vector
+Color: TypeAlias = mathutils.Color
+Euler: TypeAlias = mathutils.Euler
+Quaternion: TypeAlias = mathutils.Quaternion
+Matrix: TypeAlias = mathutils.Matrix
+BVHTree: TypeAlias = mathutils.bvhtree.BVHTree
+
+NodeGroup: TypeAlias = bpy.types.NodeGroup
+Scene: TypeAlias = bpy.types.Scene
+ViewLayer: TypeAlias = bpy.types.ViewLayer
+
+
+class Object(BlenderAsset[bpy.types.Object]):
+    def __init__(self, obj: bpy.types.Object):
+        assert isinstance(obj, bpy.types.Object)
+        super().__init__(obj)
+
+    def __repr__(self):
+        return f"pf.{self.__class__.__name__}(bpy.data.objects[{self._item.name!r}])"
+
+    def clone(self):
+        new_obj = self._item.copy()
+        new_obj.data = self._item.data.copy()
+        bpy.context.collection.objects.link(new_obj)
+        return self.__class__(new_obj)
+
+
+class CameraObject(Object):
+    def __init__(self, obj: bpy.types.Object):
+        assert obj.type == ObjectType.CAMERA.value
+        super().__init__(obj)
+
+
+class MeshObject(Object):
+    def __init__(self, obj: bpy.types.Object):
+        assert obj.type == ObjectType.MESH.value
+        super().__init__(obj)
+
+        if len(obj.children) > 0:
+            raise ValueError(
+                f"MeshObject {obj.name} had children {obj.children}, but this is not allowed for {self.__class__.__name__}"
+            )
+        if obj.parent is not None:
+            logger.warning(f"MeshObject {obj.name} had a parent {obj.parent}")
+
+
+class CurveObject(Object):
+    def __init__(self, obj: bpy.types.Object):
+        assert obj.type == ObjectType.CURVE.value
+        super().__init__(obj)
+
+
+class EmptyObject(Object):
+    def __init__(self, obj: bpy.types.Object):
+        assert obj.type == ObjectType.EMPTY.value
+        super().__init__(obj)
+
+
+class ArmatureObject(Object):
+    def __init__(self, obj: bpy.types.Object):
+        assert obj.type == ObjectType.ARMATURE.value
+        super().__init__(obj)
+
+
+class HairObject(Object):
+    def __init__(self, obj: bpy.types.Object):
+        assert obj.type == ObjectType.HAIR.value
+        super().__init__(obj)
+
+
+class LatticeObject(Object):
+    def __init__(self, obj: bpy.types.Object):
+        assert obj.type == ObjectType.LATTICE.value
+        super().__init__(obj)
+
+
+class LightObject(Object):
+    def __init__(self, obj: bpy.types.Object):
+        assert obj.type == ObjectType.LIGHT.value
+        super().__init__(obj)
+
+
+class LightProbeObject(Object):
+    def __init__(self, obj: bpy.types.Object):
+        assert obj.type == ObjectType.LIGHT_PROBE.value
+        super().__init__(obj)
+
+
+class MetaObject(Object):
+    def __init__(self, obj: bpy.types.Object):
+        assert obj.type == ObjectType.META.value
+        super().__init__(obj)
+
+
+@dataclass
+class Material:
+    surface: Any = None
+    displacement: Any = None
+    volume: Any = None
+    _bpy_material: Any = field(default=None, init=False, repr=False)
+
+    def item(self) -> bpy.types.Material:
+        if self._bpy_material is None:
+            from procfunc.nodes.execute.execute import _build_bpy_material
+
+            self._bpy_material = _build_bpy_material(
+                surface=self.surface,
+                displacement=self.displacement,
+                volume=self.volume,
+            )
+        return self._bpy_material
+
+
+class Texture(BlenderAsset[bpy.types.Texture]):
+    def __init__(self, tex: bpy.types.Texture):
+        assert isinstance(tex, bpy.types.Texture)
+        super().__init__(tex)
+
+
+class Image(BlenderAsset[bpy.types.Image]):
+    def __init__(self, img: bpy.types.Image):
+        assert isinstance(img, bpy.types.Image)
+        super().__init__(img)
+
+
+class Collection:
+    def __init__(
+        self, objects: "list[Object] | bpy.types.Collection", name: str = "collection"
+    ):
+        if isinstance(objects, bpy.types.Collection):
+            self._collection = objects
+            self._objects = [Object(o) for o in objects.objects]
+        else:
+            self._objects = objects
+            self._collection = bpy.data.collections.new(name=name)
+            for obj in objects:
+                self._collection.objects.link(obj.item())
+
+    def __repr__(self):
+        return f"pf.Collection(bpy.data.collections[{self._collection.name!r}])"
+
+    def map(
+        self,
+        fn: Callable[[Object], Object],
+        skip_none: bool = True,
+    ) -> "Collection":
+        objs = []
+        for obj in self._objects:
+            result = fn(obj)
+            if result is None and skip_none:
+                continue
+            objs.append(result)
+        return Collection(objs, name=self._collection.name + "_" + fn.__name__)
+
+    def __len__(self):
+        return len(self._objects)
+
+    def __iter__(self):
+        return iter(self._objects)
+
+    def item(self) -> bpy.types.Collection:
+        return self._collection
+
+
+class VolumeObject(BlenderAsset[bpy.types.Volume]):
+    def __init__(self, vol: bpy.types.Volume):
+        assert isinstance(vol, bpy.types.Volume)
+        super().__init__(vol)
+
+
+class PointCloudObject(BlenderAsset[bpy.types.PointCloud]):
+    def __init__(self, vol: bpy.types.PointCloud):
+        assert isinstance(vol, bpy.types.PointCloud)
+        super().__init__(vol)
+
+
+class World(BlenderAsset[bpy.types.World]):
+    def __init__(self, world: bpy.types.World):
+        assert isinstance(world, bpy.types.World)
+        super().__init__(world)
+
+    def __del__(self):
+        pass  # no cleanup for these
+
+
+# TODO Font, PointCloud, Surface, Volume, Scene? World? ViewLayer?
+
+TRangeType = TypeVar("TRangeType")
+
+
+class ValueRange(NamedTuple, Generic[TRangeType]):
+    min: TRangeType | None
+    max: TRangeType | None
+
+
+register_pytree_container(
+    Material,
+    flatten_func=lambda m: ([m.surface, m.displacement, m.volume], None),
+    unflatten_func=lambda vals, _: Material(
+        surface=vals[0], displacement=vals[1], volume=vals[2]
+    ),
+    names_func=lambda m: ["surface", "displacement", "volume"],
+)