Add support for basic feature flags

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