Repository: pygfx/pygfx
Pull request #718 (2024-04-10 13:04:58)

Refactor enums
We kindof used enums in a few places, but with the large number of marker types in #708 it was time to have a good look at enums. This also touches on #272 a bit.

* [x] Implemented an enum class. 
   * I played a bit with Python's `enum.Enum` but I hate that `SomeEnum.field` provides ... an enum object.
   * In our custom enum it simply produces a string or an int (depening on enum / flag-like).
   * This also makes things simpler when these values reach the templated shader code.
   * I started by using a similar implementation [as in wgpu](https://github.com/pygfx/wgpu-py/blob/e8bef7db15462772c048ff2f325f618046eec2f8/wgpu/enums.py#L12-L27), which works perfectly, but it takes a lot of effort to make Sphynx document them properly.
   * So I used classes instead. As a result I needed a metaclass (yuk), but it's still very little code.
* [x] Moved (I think) all  current enums to `enums.py`.
* [x] I tweaked the docs template a bit. No effect in the end, but a bit clearer, I think. 
* [ ] ~Add some logic to nicely support aliases (e.g. "+" for "plus" in markers).~ Solve in #708

It makes the docs really nice, the circles below are links:
<img width="652" alt="image" src="https://github.com/pygfx/pygfx/assets/3015475/13fa06f5-f1c0-48a5-8d5a-42d2a3ea8e41">

Follow-ups:
* Maybe use the same enum class in wgpu-py. 
* New enum for SizeMode (mode/world/screen).
* Proceed with #708.
* Probably more enums we can add ... also with #272. 


Korijn: In principle I think a custom enum implementation is a bad idea, not worth all the lines of code, just to have its members be a string...

Also, py 3.11 has [StrEnum](https://docs.python.org/3/library/enum.html#enum.StrEnum), which I guess is what you're looking for?
almarklein: I addressed your points.

Patch:
diff --git a/docs/_templates/custom_layout.rst b/docs/_templates/custom_layout.rst
index 55f848f..368075c 100644
--- a/docs/_templates/custom_layout.rst
+++ b/docs/_templates/custom_layout.rst
@@ -2,7 +2,9 @@
 
 .. currentmodule:: {{ module }}
 
+
 {% if objtype == "module" %}
+{# Assume that the module docstring has an autosummary section to produce a list of members. #}
 
 .. automodule:: {{ fullname }}
 
@@ -14,7 +16,7 @@
     :add-heading: Examples
     :heading-level: ^
 
-{% else %}
+{% elif objtype == "class" %}
 
 .. autoclass:: {{ objname }}
     :members:
@@ -24,4 +26,8 @@
 .. minigallery:: pygfx.{{ objname }}
     :add-heading: Examples
 
+{% else %}
+
+.. autodata:: {{ objname }}
+
 {% endif %}
diff --git a/pygfx/__init__.py b/pygfx/__init__.py
index 4b53cc0..591f7a2 100644
--- a/pygfx/__init__.py
+++ b/pygfx/__init__.py
@@ -17,7 +17,8 @@ from .utils.load import load_mesh, load_meshes, load_scene
 from .utils.show import show, Display
 from .utils.viewport import Viewport
 from .utils.text import font_manager
-from .utils import cm, logger
+from .utils import cm, enums, logger
+from .utils.enums import *
 
 # Temp fix for pyinstaller to pick up pylinalg
 import pylinalg
diff --git a/pygfx/materials/_base.py b/pygfx/materials/_base.py
index cd68aa9..8a0833d 100644
--- a/pygfx/materials/_base.py
+++ b/pygfx/materials/_base.py
@@ -1,5 +1,3 @@
-import enum
-
 from ..utils.trackable import Trackable
 from ..utils import array_from_shadertype
 from ..resources import Buffer
@@ -201,17 +199,3 @@ class Material(Trackable):
         if not isinstance(value, (bool, int)):
             raise TypeError("Material.pick_write must be bool.")
         self._store.pick_write = bool(value)
-
-
-class ColorMode(enum.Enum):
-    auto = 0
-    uniform = 1
-    vertex = 2
-    face = 3
-    vertex_map = 4
-    face_map = 5
-
-
-class SizeMode(enum.Enum):
-    uniform = 1
-    vertex = 2
diff --git a/pygfx/materials/_line.py b/pygfx/materials/_line.py
index 4ee57a5..15e3608 100644
--- a/pygfx/materials/_line.py
+++ b/pygfx/materials/_line.py
@@ -1,6 +1,7 @@
-from ._base import Material, ColorMode
+from ._base import Material
 from ..resources import Texture
 from ..utils import unpack_bitfield, Color
+from ..utils.enums import ColorMode
 
 
 class LineMaterial(Material):
@@ -14,7 +15,7 @@ class LineMaterial(Material):
         The coordinate space in which the thickness is expressed ('screen', 'world', 'model'). Default 'screen'.
     color : Color
         The uniform color of the line (used depending on the ``color_mode``).
-    color_mode : enum or str
+    color_mode : str | ColorMode
         The mode by which the line is coloured. Default 'auto'.
     map : Texture
         The texture map specifying the color for each texture coordinate. Optional.
@@ -115,28 +116,16 @@ class LineMaterial(Material):
     def color_mode(self):
         """The way that color is applied to the line.
 
-        * auto: switch between `uniform` and `vertex_map`, depending on whether `map` is set.
-        * uniform: use the material's color property for the whole line.
-        * vertex: use the geometry `colors` buffer, one color per vertex.
-        * face: use the geometry `colors` buffer, one color per line-piece.
-        * vertex_map: use the geometry texcoords buffer to sample (per vertex) in the material's ``map`` texture.
-        * faces_map: use the geometry texcoords buffer to sample (per line-piece) in the material's ``map`` texture.
+        See :obj:`pygfx.utils.enums.ColorMode`:
         """
         return self._store.color_mode
 
     @color_mode.setter
     def color_mode(self, value):
-        if isinstance(value, ColorMode):
-            pass
-        elif isinstance(value, str):
-            if value.startswith("ColorMode."):
-                value = value.split(".")[-1]
-            try:
-                value = getattr(ColorMode, value.lower())
-            except AttributeError:
-                raise ValueError(f"Invalid color_mode: '{value}'")
-        else:
-            raise TypeError(f"Invalid color_mode class: {value.__class__.__name__}")
+        if value not in ColorMode:
+            raise ValueError(
+                f"LineMaterial.color_mode must be a string in {ColorMode}, not {repr(value)}"
+            )
         self._store.color_mode = value
 
     @property
@@ -168,6 +157,7 @@ class LineMaterial(Material):
         """The coordinate space in which the thickness (and dash_pattern) are expressed.
 
         Possible values are:
+
         * "screen": logical screen pixels. The Default.
         * "world": the world / scene coordinate frame.
         * "model": the line's local coordinate frame (same as the line's positions).
diff --git a/pygfx/materials/_mesh.py b/pygfx/materials/_mesh.py
index 3e73fe1..5c0e2e0 100644
--- a/pygfx/materials/_mesh.py
+++ b/pygfx/materials/_mesh.py
@@ -1,8 +1,9 @@
 import math
-from ._base import Material, ColorMode
+from ._base import Material
 from ..resources import Texture
 from ..utils import logger
 from ..utils.color import Color
+from ..utils.enums import ColorMode
 
 
 class MeshAbstractMaterial(Material):
@@ -14,7 +15,7 @@ class MeshAbstractMaterial(Material):
     ----------
     color : Color
         The uniform color of the mesh (used depending on the ``color_mode``).
-    color_mode : enum or str
+    color_mode : str | ColorMode
         The mode by which the mesh is coloured. Default 'auto'.
     map : Texture
         The texture map specifying the color at each texture coordinate. Optional.
@@ -89,29 +90,17 @@ class MeshAbstractMaterial(Material):
     def color_mode(self):
         """The way that color is applied to the mesh.
 
-        * auto: switch between `uniform` and `vertex_map`, depending on whether `map` is set.
-        * uniform: use the material's color property for the whole mesh.
-        * vertex: use the geometry `colors` buffer, one color per vertex.
-        * face: use the geometry `colors` buffer, one color per face.
-        * vertex_map: use the geometry texcoords buffer to sample (per vertex) in the material's ``map`` texture.
-        * faces_map: use the geometry texcoords buffer to sample (per face) in the material's ``map`` texture.
+        See :obj:`pygfx.utils.enums.ColorMode`:
         """
         # todo: does 'auto' take the presence of texcoords into account?
         return self._store.color_mode
 
     @color_mode.setter
     def color_mode(self, value):
-        if isinstance(value, ColorMode):
-            pass
-        elif isinstance(value, str):
-            if value.startswith("ColorMode."):
-                value = value.split(".")[-1]
-            try:
-                value = getattr(ColorMode, value.lower())
-            except AttributeError:
-                raise ValueError(f"Invalid color_mode: '{value}'")
-        else:
-            raise TypeError(f"Invalid color_mode class: {value.__class__.__name__}")
+        if value not in ColorMode:
+            raise ValueError(
+                f"MeshMaterial.color_mode must be a string in {ColorMode}, not {repr(value)}"
+            )
         self._store.color_mode = value
 
     @property
diff --git a/pygfx/materials/_points.py b/pygfx/materials/_points.py
index 9d8ee82..acfc9f2 100644
--- a/pygfx/materials/_points.py
+++ b/pygfx/materials/_points.py
@@ -1,6 +1,7 @@
-from ._base import Material, ColorMode, SizeMode
+from ._base import Material
 from ..resources import Texture
 from ..utils import unpack_bitfield, Color
+from ..utils.enums import ColorMode, SizeMode
 
 
 class PointsMaterial(Material):
@@ -14,11 +15,11 @@ class PointsMaterial(Material):
         The size (diameter) of the points in logical pixels. Default 4.
     size_space : str
         The coordinate space in which the size is expressed ('screen', 'world', 'model'). Default 'screen'.
-    size_mode : enum or str
+    size_mode : str | SizeMode
         The mode by which the points are sized. Default 'uniform'.
-    color : Color
+    color : str | tuple | Color
         The uniform color of the points (used depending on the ``color_mode``).
-    color_mode : enum or str
+    color_mode : str | ColorMode
         The mode by which the points are coloured. Default 'auto'.
     map : Texture
         The texture map specifying the color for each texture coordinate.
@@ -114,28 +115,16 @@ class PointsMaterial(Material):
     def color_mode(self):
         """The way that color is applied to the mesh.
 
-        * auto: switch between `uniform` and `vertex_map`, depending on whether `map` is set.
-        * uniform: use the material's color property for the whole mesh.
-        * vertex: use the geometry `colors` buffer, one color per vertex.
-        * vertex_map: use the geometry texcoords buffer to sample (per vertex) in the material's ``map`` texture.
+        See :obj:`pygfx.utils.enums.ColorMode`:
         """
         return self._store.color_mode
 
     @color_mode.setter
     def color_mode(self, value):
-        if isinstance(value, ColorMode):
-            pass
-        elif isinstance(value, str):
-            if value.startswith("ColorMode."):
-                value = value.split(".")[-1]
-            try:
-                value = getattr(ColorMode, value.lower())
-            except AttributeError:
-                raise ValueError(f"Invalid color_mode: '{value}'")
-        else:
-            raise TypeError(f"Invalid color_mode class: {value.__class__.__name__}")
-        if value == ColorMode.face or value == ColorMode.face_map:
-            raise ValueError(f"Points cannot have color_mode {value}")
+        if value not in ColorMode:
+            raise ValueError(
+                f"PointsMaterial.color_mode must be a string in {ColorMode}, not {repr(value)}"
+            )
         self._store.color_mode = value
 
     @property
@@ -163,6 +152,7 @@ class PointsMaterial(Material):
         """The coordinate space in which the size is expressed.
 
         Possible values are:
+
         * "screen": logical screen pixels. The Default.
         * "world": the world / scene coordinate frame.
         * "model": the line's local coordinate frame (same as the line's positions).
@@ -184,24 +174,16 @@ class PointsMaterial(Material):
     def size_mode(self):
         """The way that size is applied to the mesh.
 
-        * uniform: use the material's size property for all points.
-        * vertex: use the geometry `sizes` buffer, one size per vertex.
+        See :obj:`pygfx.utils.enums.SizeMode`:
         """
         return self._store.size_mode
 
     @size_mode.setter
     def size_mode(self, value):
-        if isinstance(value, SizeMode):
-            pass
-        elif isinstance(value, str):
-            if value.startswith("SizeMode."):
-                value = value.split(".")[-1]
-            try:
-                value = getattr(SizeMode, value.lower())
-            except AttributeError:
-                raise ValueError(f"Invalid size_mode: '{value}'")
-        else:
-            raise TypeError(f"Invalid size_mode class: {value.__class__.__name__}")
+        if value not in SizeMode:
+            raise ValueError(
+                f"PointsMaterial.size_mode must be a string in {SizeMode}, not {repr(value)}"
+            )
         self._store.size_mode = value
 
     @property
diff --git a/pygfx/objects/_base.py b/pygfx/objects/_base.py
index b6e3bb9..3854a59 100644
--- a/pygfx/objects/_base.py
+++ b/pygfx/objects/_base.py
@@ -1,7 +1,6 @@
 import random
 import weakref
 import threading
-import enum
 from typing import List
 import pylinalg as la
 
@@ -17,6 +16,7 @@ from ..utils.transform import (
     RecursiveTransform,
     callback,
 )
+from ..utils.enums import RenderMask
 
 
 class IdProvider:
@@ -75,20 +75,6 @@ class IdProvider:
 id_provider = IdProvider()
 
 
-class RenderMask(enum.IntFlag):
-    """A flag indicating how the object participates in the different render passes.
-
-    It defines whether an object should be rendered in the opaque pass, the
-    transparent pass, or all. By default (auto), this is determined from the
-    materials properties (e.g. opacity, color, color_mode).
-    """
-
-    auto = 0
-    opaque = 1
-    transparent = 2
-    all = 3
-
-
 class WorldObject(EventTarget, RootTrackable):
     """Base class for objects.
 
@@ -108,7 +94,7 @@ class WorldObject(EventTarget, RootTrackable):
         Whether the object is visible.
     render_order : float
         Value that helps controls the order in which objects are rendered.
-    render_mask : str
+    render_mask : str | RenderMask
         Determines the render passes that the object is rendered in. It's
         recommended to let the renderer decide, using "auto".
 
@@ -257,10 +243,7 @@ class WorldObject(EventTarget, RootTrackable):
     def render_mask(self):
         """Indicates in what render passes to render this object:
 
-        * "auto": try to determine the best approach (default).
-        * "opaque": only in the opaque render pass.
-        * "transparent": only in the transparent render pass(es).
-        * "all": render in both opaque and transparent render passses.
+        See :obj:`pygfx.utils.enums.RenderMask`:
 
         If "auto" (the default), the renderer attempts to determine
         whether all fragments will be either opaque or all transparent,
@@ -283,20 +266,20 @@ class WorldObject(EventTarget, RootTrackable):
     @render_mask.setter
     def render_mask(self, value):
         if value is None:
-            self._store.render_mask = RenderMask(0)
-        elif isinstance(value, int):
-            self._store.render_mask = RenderMask(value)
+            value = 0
+        if isinstance(value, int):
+            pass
         elif isinstance(value, str):
-            try:
-                self._store.render_mask = RenderMask._member_map_[value.lower()]
-            except KeyError:
-                opts = set(RenderMask._member_names_)
-                msg = f"WorldObject.render_mask must be one of {opts} not {value!r}"
+            if value not in dir(RenderMask):
+                msg = f"WorldObject.render_mask must be one of {dir(RenderMask)} not {value!r}"
                 raise ValueError(msg) from None
+            value = RenderMask[value]
         else:
             raise TypeError(
                 f"WorldObject.render_mask must be int or str, not {type(value)}"
             )
+        # Store the value as an int, because this is a flag, but also for backwards compat.
+        self._store.render_mask = value
 
     @property
     def geometry(self):
diff --git a/pygfx/renderers/wgpu/__init__.py b/pygfx/renderers/wgpu/__init__.py
index a6fb450..6b5fae4 100644
--- a/pygfx/renderers/wgpu/__init__.py
+++ b/pygfx/renderers/wgpu/__init__.py
@@ -25,7 +25,6 @@ Classes and functions required to implement custom shaders:
     :template: ../_templates/custom_layout.rst
 
     WorldObjectShader
-    RenderMask
     Binding
     GfxSampler
     GfxTextureView
@@ -50,7 +49,7 @@ Lower level functions that may or may not be needed in custom shaders:
 # TODO: \ register_wgpu_render_function to register_wgpu_shader
 
 # Import stuff that people need who create custom shaders, so they can import from pygfx.renderers.wgpu
-from ...objects._base import RenderMask
+from ...utils.enums import RenderMask
 from .engine.utils import (
     register_wgpu_render_function,
     nchannels_from_format,
diff --git a/pygfx/utils/__init__.py b/pygfx/utils/__init__.py
index 4a260b7..8d6dffa 100644
--- a/pygfx/utils/__init__.py
+++ b/pygfx/utils/__init__.py
@@ -15,6 +15,7 @@ Utility functions for pygfx.
     show.Display
     viewport.Viewport
     text.FontManager
+    enums
     cm
 
 Transform classes
@@ -40,6 +41,7 @@ import logging
 import numpy as np
 
 from .color import Color  # noqa: F401
+from . import enums  # noqa: F401
 from . import cm  # noqa: F401
 from ._dirs import get_resources_dir, get_cache_dir  # noqa: F401
 
diff --git a/pygfx/utils/enums.py b/pygfx/utils/enums.py
new file mode 100644
index 0000000..043e081
--- /dev/null
+++ b/pygfx/utils/enums.py
@@ -0,0 +1,114 @@
+"""
+The enums used in pygfx. The enums are all available from the root ``pygfx`` namespace.
+
+.. currentmodule:: pygfx.utils.enums
+
+.. autosummary::
+    :toctree: utils/enums
+    :template: ../_templates/custom_layout.rst
+
+    RenderMask
+    ColorMode
+    SizeMode
+
+"""
+
+__all__ = ["RenderMask", "ColorMode", "SizeMode"]
+
+# We implement a custom enum class that's much simpler than Python's enum.Enum,
+# and simply maps to strings or ints. The enums are classes, so IDE's provide
+# autocompletion, and documenting with Sphinx is easy. That does mean we need a
+# metaclass though.
+
+import types
+
+
+class EnumType(type):
+    """Enum metaclas."""
+
+    def __new__(cls, name, bases, dct):
+        # Collect and check fields
+        member_map = {}
+        for key, val in dct.items():
+            if not key.startswith("_"):
+                val = key if val is None else val
+                if not isinstance(val, (int, str)):
+                    raise TypeError("Enum fields must be str or int.")
+                member_map[key] = val
+        # Some field values may have been updated
+        dct.update(member_map)
+        # Create class
+        klass = super().__new__(cls, name, bases, dct)
+        # Attach some fields
+        klass.__fields__ = tuple(member_map)
+        klass.__members__ = types.MappingProxyType(member_map)  # enums.Enum compat
+        # Create bound methods
+        for name in ["__dir__", "__iter__", "__getitem__", "__setattr__", "__repr__"]:
+            setattr(klass, name, types.MethodType(getattr(cls, name), klass))
+        return klass
+
+    def __dir__(cls):
+        # Support dir(enum). Note that this order matches the definition, but dir() makes it alphabetic.
+        return cls.__fields__
+
+    def __iter__(cls):
+        # Support list(enum), iterating over the enum, and doing ``x in enum``.
+        return iter([getattr(cls, key) for key in cls.__fields__])
+
+    def __getitem__(cls, key):
+        # Support enum[key]
+        return cls.__dict__[key]
+
+    def __repr__(cls):
+        name = cls.__name__
+        options = []
+        for key in cls.__fields__:
+            val = cls[key]
+            options.append(f"'{key}' ({val})" if isinstance(val, int) else f"'{val}'")
+        return f"<pygfx.{name} enum with options: {', '.join(options)}>"
+
+    def __setattr__(cls, name, value):
+        if name.startswith("_"):
+            super().__setattr__(name, value)
+        else:
+            raise RuntimeError("Cannot set values on an enum.")
+
+
+class Enum(metaclass=EnumType):
+    """Enum base class."""
+
+    def __init__(self):
+        raise RuntimeError("Connot instantiate an enum.")
+
+
+# --- The enums
+
+
+class RenderMask(Enum):
+    """The RenderMask enum specifies the render passes in which an object participates."""
+
+    auto = 0  #: Select the appropriate render passes automatically.
+    opaque = 1  #: Only render in the opaque pass.
+    transparent = 2  #: Only render in the transparancy pass.
+    all = 3  #: Render in both passes.
+
+
+class ColorMode(Enum):
+    """The ColorMode enum specifies how an object's color is established."""
+
+    auto = None  #: Use either ``uniform`` and ``vertex_map``, depending on whether ``map`` is set.
+    uniform = None  #: Use the uniform color (usually ``material.color``).
+    vertex = None  #: Use the per-vertex color specified in the geometry  (usually  ``geometry.colors``).
+    face = None  #: Use the per-face color specified in the geometry  (usually  ``geometry.colors``).
+    vertex_map = None  #: Use per-vertex texture coords (``geometry.texcoords``), and sample these in ``material.map``.
+    face_map = None  #: Use per-face texture coords (``geometry.texcoords``), and sample these in ``material.map``.
+
+
+class SizeMode(Enum):
+    """The SizeMode enum specifies how an object's size/width/thickness is established."""
+
+    uniform = None  #: Use a uniform size.
+    vertex = None  #: Use a per-vertex size specified on the geometry.
+
+
+# NOTE: Don't forget to add new enums to the toctree and __all__


Test patch:
diff --git a/tests/utils/test_enums.py b/tests/utils/test_enums.py
new file mode 100644
index 0000000..b59a694
--- /dev/null
+++ b/tests/utils/test_enums.py
@@ -0,0 +1,59 @@
+from pygfx.utils.enums import Enum
+from pytest import raises
+
+
+def test_enums():
+    class MyOption(Enum):
+        auto = "auto"  # fields map to str or int
+        some_attr = "some-attr"  # wgpu-style values
+        foo = None  # value is the same as the key, most-used in pygfx
+
+    # Use dir() to get an (alphabetic) list of keys / options.
+    assert dir(MyOption) == ["auto", "foo", "some_attr"]
+
+    # Iterate over the object to get a list of values, in original order.
+    assert list(MyOption) == ["auto", "some-attr", "foo"]
+
+    # The repr is actually useful.
+    assert (
+        str(MyOption)
+        == "<pygfx.MyOption enum with options: 'auto', 'some-attr', 'foo'>"
+    )
+
+    # Attribute and map-like lookups are supported
+    assert MyOption.some_attr == "some-attr"
+    assert MyOption["some_attr"] == "some-attr"
+
+    # Enums are 'immutable'
+    with raises(RuntimeError):
+        MyOption.auto = "foo"
+
+
+def test_flags():
+
+    # Our Enum class does flags too.
+
+    class MyFlag(Enum):
+        auto = 0
+        foo = 1
+        bar = 2
+
+    # Use dir() to get an (alphabetic) list of keys / options.
+    assert dir(MyFlag) == ["auto", "bar", "foo"]
+
+    # Iterate over the object to get a list of values, in original order.
+    assert list(MyFlag) == [0, 1, 2]
+
+    # The repr is actually useful.
+    assert (
+        str(MyFlag)
+        == "<pygfx.MyFlag enum with options: 'auto' (0), 'foo' (1), 'bar' (2)>"
+    )
+
+    # Attribute and map-like lookups are supported
+    assert MyFlag.bar == 2
+    assert MyFlag["bar"] == 2
+
+    # Enums are 'immutable'
+    with raises(RuntimeError):
+        MyFlag.auto = 99