Merge pull request #83 from Carifio24/registries

Use registries

glue_ar/__init__.py

@@ -6,6 +6,14 @@
     pass
 
 
+def setup_common():
+    from .common.gltf_builder import GLTFBuilder  # noqa: F401
+    from .common.usd_builder import USDBuilder  # noqa: F401
+    from .common.stl_builder import STLBuilder  # noqa: F401
+
+    from .compression import compress_gltfpack, compress_gltf_pipeline  # noqa: F401
+
+
 def setup_qt():
     try:
         from glue_vispy_viewers.scatter.qt.scatter_viewer import VispyScatterViewer
@@ -56,6 +64,9 @@ def setup_jupyter():
 
 
 def setup():
+
+    setup_common()
+
     try:
         setup_qt()
     except ImportError:

glue_ar/common/__init__.py

@@ -1,3 +1,6 @@
+from .gltf_builder import GLTFBuilder  # noqa: F401
+from .usd_builder import USDBuilder  # noqa: F401
+from .stl_builder import STLBuilder  # noqa: F401
 from .marching_cubes import add_isosurface_layer_gltf, add_isosurface_layer_usd  # noqa: F401
 from .scatter_gltf import add_scatter_layer_gltf  # noqa: F401
 from .scatter_stl import add_scatter_layer_stl  # noqa: F401

glue_ar/common/export.py

@@ -2,7 +2,6 @@
 from inspect import getfullargspec
 from os.path import extsep, join, split, splitext
 from string import Template
-from subprocess import run
 from typing import Dict, Optional
 from glue.core.state_objects import State
 from glue.config import settings
@@ -11,29 +10,14 @@
 
 
 from glue_ar.common.export_options import ar_layer_export
-from glue_ar.common.gltf_builder import GLTFBuilder
-from glue_ar.common.stl_builder import STLBuilder
-from glue_ar.common.usd_builder import USDBuilder
+from glue_ar.registries import builder as builder_registry, compressor as compressor_registry
 from glue_ar.utils import PACKAGE_DIR, RESOURCES_DIR, Bounds, BoundsWithResolution, export_label_for_layer
 
 from typing import List, Tuple, Union
 
 
 NODE_MODULES_DIR = join(PACKAGE_DIR, "js", "node_modules")
 
-GLTF_PIPELINE_FILEPATH = join(NODE_MODULES_DIR, "gltf-pipeline", "bin", "gltf-pipeline.js")
-GLTFPACK_FILEPATH = join(NODE_MODULES_DIR, "gltfpack", "cli.js")
-
-
-_BUILDERS = {
-    "gltf": GLTFBuilder,
-    "glb": GLTFBuilder,
-    "usda": USDBuilder,
-    "usdc": USDBuilder,
-    "usdz": USDBuilder,
-    "stl": STLBuilder,
-}
-
 
 def export_viewer(viewer_state: Vispy3DViewerState,
                   layer_states: List[VolumeLayerState],
@@ -45,7 +29,7 @@ def export_viewer(viewer_state: Vispy3DViewerState,
 
     base, ext = splitext(filepath)
     ext = ext[1:]
-    builder_cls = _BUILDERS[ext]
+    builder_cls = builder_registry.members.get(ext)
     count = len(getfullargspec(builder_cls.__init__)[0])
     builder_args = [filepath] if count > 1 else []
     builder = builder_cls(*builder_args)
@@ -77,22 +61,8 @@ def export_viewer(viewer_state: Vispy3DViewerState,
             export_modelviewer(mv_path, filepath, viewer_state.title)
 
 
-def compress_gltf_pipeline(filepath: str):
-    run(["node", GLTF_PIPELINE_FILEPATH, "-i", filepath, "-o", filepath, "-d"], capture_output=True)
-
-
-def compress_gltfpack(filepath: str):
-    run(["node", GLTFPACK_FILEPATH, "-i", filepath, "-o", filepath], capture_output=True)
-
-
-COMPRESSORS = {
-    "draco": compress_gltf_pipeline,
-    "meshoptimizer": compress_gltfpack
-}
-
-
 def compress_gl(filepath: str, method: str = "draco"):
-    compressor = COMPRESSORS.get(method.lower(), None)
+    compressor = compressor_registry.members.get(method.lower(), None)
     if compressor is None:
         raise ValueError("Invalid compression method specified")
     compressor(filepath)

glue_ar/common/gltf_builder.py

@@ -7,7 +7,10 @@
 from gltflib.gltf_resource import FileResource
 from typing import Iterable, List, Optional, Union
 
+from glue_ar.registries import builder
 
+
+@builder(("gltf", "glb"))
 class GLTFBuilder:
 
     def __init__(self):

glue_ar/common/stl_builder.py

@@ -4,7 +4,10 @@
 from stl import Mesh
 from typing import Iterable, List
 
+from glue_ar.registries import builder
 
+
+@builder("stl")
 class STLBuilder:
 
     def __init__(self):

glue_ar/common/usd_builder.py

@@ -5,13 +5,15 @@
 from pxr import Usd, UsdGeom, UsdLux, UsdShade, UsdUtils
 from typing import Dict, Iterable, Optional, Tuple
 
+from glue_ar.registries import builder
 from glue_ar.usd_utils import material_for_color, material_for_mesh
 from glue_ar.utils import unique_id
 
 
 MaterialInfo = Tuple[int, int, int, float, float, float]
 
 
+@builder(("usda", "usdc", "usdz"))
 class USDBuilder:
 
     def __init__(self, filepath: str):

glue_ar/compression.py

@@ -0,0 +1,20 @@
+from os.path import join
+from subprocess import run
+
+from glue_ar.registries import compressor
+from glue_ar.utils import PACKAGE_DIR
+
+
+NODE_MODULES_DIR = join(PACKAGE_DIR, "js", "node_modules")
+GLTF_PIPELINE_FILEPATH = join(NODE_MODULES_DIR, "gltf-pipeline", "bin", "gltf-pipeline.js")
+GLTFPACK_FILEPATH = join(NODE_MODULES_DIR, "gltfpack", "cli.js")
+
+
+@compressor("draco")
+def compress_gltf_pipeline(filepath: str):
+    run(["node", GLTF_PIPELINE_FILEPATH, "-i", filepath, "-o", filepath, "-d"], capture_output=True)
+
+
+@compressor("meshoptimizer")
+def compress_gltfpack(filepath: str):
+    run(["node", GLTFPACK_FILEPATH, "-i", filepath, "-o", filepath], capture_output=True)

glue_ar/registries.py

@@ -0,0 +1,41 @@
+from collections.abc import Callable
+from typing import Tuple, Type, Union
+
+from glue.config import DictRegistry
+
+
+__all__ = ["builder", "compressor"]
+
+
+class BuilderRegistry(DictRegistry):
+
+    def add(self, extensions: Union[str, Tuple[str]], builder: Type):
+        if isinstance(extensions, str):
+            self._members[extensions] = builder
+        else:
+            for ext in extensions:
+                self._members[ext] = builder
+
+    def __call__(self, extensions: Union[str, Tuple[str]]):
+        def adder(builder: Type):
+            self.add(extensions, builder)
+            return builder
+        return adder
+
+
+builder = BuilderRegistry()
+
+
+class CompressorRegistry(DictRegistry):
+
+    def add(self, name: str, compressor: Callable[[str], None]):
+        self._members[name] = compressor
+
+    def __call__(self, name: str):
+        def adder(compressor: Callable[[str], None]):
+            self.add(name, compressor)
+            return compressor
+        return adder
+
+
+compressor = CompressorRegistry()

glue_ar/tests/test_registries.py

@@ -0,0 +1,29 @@
+from ..registries import BuilderRegistry, CompressorRegistry
+
+
+def test_builder_registry():
+
+    builder_registry = BuilderRegistry()
+
+    @builder_registry("ext")
+    class ExtBuilder:
+        pass
+
+    @builder_registry(("ext1", "ext2"))
+    class ExtsBuilder:
+        pass
+
+    assert builder_registry.members == {"ext": ExtBuilder,
+                                        "ext1": ExtsBuilder,
+                                        "ext2": ExtsBuilder}
+
+
+def test_compressor_registry():
+
+    compressor_registry = CompressorRegistry()
+
+    @compressor_registry("dummy")
+    def dummy_compressor(_filepath: str):
+        pass
+
+    assert compressor_registry.members == {"dummy": dummy_compressor}