Add support for basic feature flags (#35)

(DIS-1356)

dissect/util/feature.py

@@ -0,0 +1,74 @@
+import functools
+import os
+from enum import Enum
+from typing import Callable, Optional
+
+
+# Register feature flags in a central place to avoid chaos
+class Feature(Enum):
+    LATEST = "latest"
+    ADVANCED = "advanced"
+    BETA = "beta"
+
+
+# Defines the default flags (as strings)
+DISSECT_FEATURES_DEFAULT = "latest"
+
+# Defines the environment variable to read the flags from
+DISSECT_FEATURES_ENV = "DISSECT_FEATURES"
+
+
+class FeatureException(RuntimeError):
+    pass
+
+
+@functools.cache
+def feature_flags() -> list[Feature]:
+    return [Feature(name) for name in os.getenv(DISSECT_FEATURES_ENV, DISSECT_FEATURES_DEFAULT).split("/")]
+
+
+@functools.cache
+def feature_enabled(feature: Feature) -> bool:
+    """Use this function for block-level feature flag control.
+
+    Usage::
+
+        def parse_blob():
+            if feature_enabled(Feature.BETA):
+                self._parse_fast_experimental()
+            else:
+                self._parse_normal()
+
+    """
+    return feature in feature_flags()
+
+
+def feature(flag: Feature, alternative: Optional[Callable] = None) -> Callable:
+    """Feature flag decorator allowing you to guard a function behind a feature flag.
+
+    Usage::
+
+        @feature(Feature.SOME_FLAG, fallback)
+        def my_func( ... ) -> ...
+
+    Where ``SOME_FLAG`` is the feature you want to check for and ``fallback`` is the alternative function to serve
+    if the feature flag is NOT set.
+    """
+
+    if alternative is None:
+
+        def alternative():
+            raise FeatureException(
+                "\n".join(
+                    [
+                        "Feature disabled.",
+                        f"Set FLAG '{flag}' in {DISSECT_FEATURES_ENV} to enable.",
+                        "See https://docs.dissect.tools/en/latest/advanced/flags.html",
+                    ]
+                )
+            )
+
+    def decorator(func):
+        return func if feature_enabled(flag) else alternative
+
+    return decorator

tests/test_feature.py

@@ -0,0 +1,35 @@
+import pytest
+
+from dissect.util.feature import Feature, FeatureException, feature, feature_enabled
+
+
+def test_feature_flags() -> None:
+    def fallback():
+        return False
+
+    @feature(Feature.BETA, fallback)
+    def experimental():
+        return True
+
+    @feature(Feature.ADVANCED, fallback)
+    def advanced():
+        return True
+
+    @feature(Feature.LATEST)
+    def latest():
+        return True
+
+    @feature("expert")
+    def expert():
+        return True
+
+    assert experimental() is False
+    assert advanced() is False
+    assert latest() is True
+    with pytest.raises(FeatureException):
+        assert expert() is True
+
+
+def test_feature_flag_inline() -> None:
+    assert feature_enabled(Feature.BETA) is False
+    assert feature_enabled(Feature.LATEST) is True