Add support for basic feature flags (DIS-1356)

dissect/util/feature.py

@@ -0,0 +1,66 @@
+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):
+    NOVICE = "novice"
+    LATEST = "latest"
+    BETA = "beta"
+
+
+# This set defines the valid flags
+DISSECT_FEATURE_SET = set(item for item in Feature)
+
+# Defines the default flags (as strings)
+DISSECT_FEATURES_DEFAULT = f"novice/latest"
+
+# Defines the environment variable to read the flags from
+DISSECT_FEATURES_ENV = "DISSECT_FEATURES"
+
+
+class FeatureException(RuntimeError):
+    pass
+
+
+def check_flags(flags: list[str]) -> list[str]:
+    for flag in flags:
+        if flag not in DISSECT_FEATURE_SET:
+            raise FeatureException(f"Invalid feature flag: {flag} choose from: {DISSECT_FEATURE_SET}")
+    return flags
+
+
+@functools.cache
+def feature_flags() -> list[str]:
+    return check_flags([Feature(name) for name in os.getenv(DISSECT_FEATURES_ENV, DISSECT_FEATURES_DEFAULT).split("/")])
+
+
+@functools.cache
+def feature_enabled(feature: Feature) -> bool:
+    return feature in feature_flags()
+
+
+def feature_disabled_stub() -> None:
+    raise FeatureException("This feature has been disabled.")
+
+
+def feature(flag: Feature, alternative: Optional[Callable] = feature_disabled_stub) -> Callable:
+    """Usage:
+
+    @feature(F_SOME_FLAG, altfunc)
+    def my_func( ... ) -> ...
+
+    Where F_SOME_FLAG is the feature you want to check for and
+    altfunc is the alternative function to serve
+    if the feature flag is NOT set.
+    """
+
+    def decorator(func):
+        if feature_enabled(flag):
+            return func
+        else:
+            return alternative
+
+    return decorator

tests/test_feature.py

@@ -0,0 +1,46 @@
+import pytest
+
+from dissect.util.feature import (
+    Feature,
+    FeatureException,
+    check_flags,
+    feature,
+    feature_enabled,
+)
+
+
+def test_feature_flags() -> None:
+    def fallback():
+        return False
+
+    @feature(Feature.BETA, fallback)
+    def experimental():
+        return True
+
+    @feature(Feature.NOVICE, fallback)
+    def novice():
+        return True
+
+    @feature(Feature.LATEST)
+    def latest():
+        return True
+
+    @feature("expert")
+    def expert():
+        return True
+
+    assert experimental() is False
+    assert novice() is True
+    assert latest() is True
+    with pytest.raises(FeatureException):
+        assert expert() is True
+
+
+def test_feature_flag_verification() -> None:
+    with pytest.raises(FeatureException):
+        check_flags(["chaotic"])
+
+
+def test_feature_flag_inline() -> None:
+    assert feature_enabled(Feature.BETA) is False
+    assert feature_enabled(Feature.LATEST) is True