Refactor default flag logic (#22)

* Refactor default flags to use a handler function instead of list

* Use default handler if API request fails

* Add more docstring

* Added gitignore for direnv

Co-authored-by: Ben Rometsch <ben.rometsch@gmail.com>

Author: Matthew Elwell

.gitignore

@@ -1,5 +1,6 @@
 .idea/
 .venv
+.direnv/
 
 *.pyc
 

.isort.cfg

@@ -3,4 +3,4 @@ use_parentheses=true
 multi_line_output=3
 include_trailing_comma=true
 line_length=79
-known_third_party = flag_engine,flask,pytest,requests,requests_futures,responses
+known_third_party = flag_engine,flask,pytest,requests,requests_futures,responses,urllib3

example/app.py

@@ -8,17 +8,25 @@
 
 app = Flask(__name__)
 
-flagsmith = Flagsmith(
-    environment_key=os.environ.get("FLAGSMITH_ENVIRONMENT_KEY"),
-    defaults=[
-        # Set a default flag which will be used if the "secret_button"
-        # feature is not returned by the API
-        DefaultFlag(
+
+def default_flag_handler(feature_name: str) -> DefaultFlag:
+    """
+    Function that will be used if the API doesn't respond, or an unknown
+    feature is requested
+    """
+
+    if feature_name == "secret_button":
+        return DefaultFlag(
             enabled=False,
             value=json.dumps({"colour": "#b8b8b8"}),
-            feature_name="secret_button",
         )
-    ],
+
+    return DefaultFlag(False, None)
+
+
+flagsmith = Flagsmith(
+    environment_key=os.environ.get("FLAGSMITH_ENVIRONMENT_KEY"),
+    default_flag_handler=default_flag_handler,
 )
 
 

flagsmith/flagsmith.py

@@ -7,7 +7,8 @@
 from flag_engine.environments.builders import build_environment_model
 from flag_engine.environments.models import EnvironmentModel
 from flag_engine.identities.models import IdentityModel, TraitModel
-from requests.adapters import HTTPAdapter, Retry
+from requests.adapters import HTTPAdapter
+from urllib3 import Retry
 
 from flagsmith.analytics import AnalyticsProcessor
 from flagsmith.exceptions import FlagsmithAPIError, FlagsmithClientError
@@ -21,6 +22,20 @@
 
 
 class Flagsmith:
+    """A Flagsmith client.
+
+    Provides an interface for interacting with the Flagsmith http API.
+
+    Basic Usage::
+
+      >>> from flagsmith import Flagsmith
+      >>> flagsmith = Flagsmith(environment_key="<your API key>")
+      >>> environment_flags = flagsmith.get_environment_flags()
+      >>> feature_enabled = environment_flags.is_feature_enabled("foo")
+      >>> identity_flags = flagsmith.get_identity_flags("identifier", {"foo": "bar"})
+      >>> feature_enabled_for_identity = identity_flags.is_feature_enabled("foo")
+    """
+
     def __init__(
         self,
         environment_key: str,
@@ -31,8 +46,26 @@ def __init__(
         environment_refresh_interval_seconds: int = 60,
         retries: Retry = None,
         enable_analytics: bool = False,
-        defaults: typing.List[DefaultFlag] = None,
+        default_flag_handler: typing.Callable[[str], DefaultFlag] = None,
     ):
+        """
+        :param environment_key: The environment key obtained from Flagsmith interface
+        :param api_url: Override the URL of the Flagsmith API to communicate with
+        :param custom_headers: Additional headers to add to requests made to the
+            Flagsmith API
+        :param request_timeout_seconds: Number of seconds to wait for a request to
+            complete before terminating the request
+        :param enable_local_evaluation: Enables local evaluation of flags
+        :param environment_refresh_interval_seconds: If using local evaluation,
+            specify the interval period between refreshes of local environment data
+        :param retries: a urllib3.Retry object to use on all http requests to the
+            Flagsmith API
+        :param enable_analytics: if enabled, sends additional requests to the Flagsmith
+            API to power flag analytics charts
+        :param default_flag_handler: callable which will be used in the case where
+            flags cannot be retrieved from the API or a non existent feature is
+            requested
+        """
         self.session = requests.Session()
         self.session.headers.update(
             **{"X-Environment-Key": environment_key}, **(custom_headers or {})
@@ -65,7 +98,7 @@ def __init__(
             else None
         )
 
-        self.defaults = defaults or []
+        self.default_flag_handler = default_flag_handler
 
     def get_environment_flags(self) -> Flags:
         """
@@ -107,7 +140,7 @@ def _get_environment_flags_from_document(self) -> Flags:
         return Flags.from_feature_state_models(
             feature_states=engine.get_environment_feature_states(self._environment),
             analytics_processor=self._analytics_processor,
-            defaults=self.defaults,
+            default_flag_handler=self.default_flag_handler,
         )
 
     def _get_identity_flags_from_document(
@@ -121,32 +154,41 @@ def _get_identity_flags_from_document(
             feature_states=feature_states,
             analytics_processor=self._analytics_processor,
             identity_id=identity_model.composite_key,
-            defaults=self.defaults,
+            default_flag_handler=self.default_flag_handler,
         )
 
     def _get_environment_flags_from_api(self) -> Flags:
-        api_flags = self._get_json_response(
-            url=self.environment_flags_url, method="GET"
-        )
-
-        return Flags.from_api_flags(
-            api_flags=api_flags,
-            analytics_processor=self._analytics_processor,
-            defaults=self.defaults,
-        )
+        try:
+            api_flags = self._get_json_response(
+                url=self.environment_flags_url, method="GET"
+            )
+            return Flags.from_api_flags(
+                api_flags=api_flags,
+                analytics_processor=self._analytics_processor,
+                default_flag_handler=self.default_flag_handler,
+            )
+        except FlagsmithAPIError:
+            if self.default_flag_handler:
+                return Flags(default_flag_handler=self.default_flag_handler)
+            raise
 
     def _get_identity_flags_from_api(
         self, identifier: str, traits: typing.Dict[str, typing.Any]
     ) -> Flags:
-        data = generate_identities_data(identifier, traits)
-        json_response = self._get_json_response(
-            url=self.identities_url, method="POST", body=data
-        )
-        return Flags.from_api_flags(
-            api_flags=json_response["flags"],
-            analytics_processor=self._analytics_processor,
-            defaults=self.defaults,
-        )
+        try:
+            data = generate_identities_data(identifier, traits)
+            json_response = self._get_json_response(
+                url=self.identities_url, method="POST", body=data
+            )
+            return Flags.from_api_flags(
+                api_flags=json_response["flags"],
+                analytics_processor=self._analytics_processor,
+                default_flag_handler=self.default_flag_handler,
+            )
+        except FlagsmithAPIError:
+            if self.default_flag_handler:
+                return Flags(default_flag_handler=self.default_flag_handler)
+            raise
 
     def _get_json_response(self, url: str, method: str, body: dict = None):
         try:

flagsmith/models.py

@@ -1,5 +1,5 @@
 import typing
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from flag_engine.features.models import FeatureStateModel
 
@@ -10,19 +10,20 @@
 @dataclass
 class BaseFlag:
     enabled: bool
-    value: typing.Union[str, int, float, bool]
-    feature_name: str
+    value: typing.Union[str, int, float, bool, type(None)]
+    is_default: bool
 
 
-@dataclass
 class DefaultFlag(BaseFlag):
-    is_default = True
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, is_default=True, **kwargs)
 
 
-@dataclass
 class Flag(BaseFlag):
-    feature_id: int
-    is_default = False
+    def __init__(self, *args, feature_id: int, feature_name: str, **kwargs):
+        super().__init__(*args, is_default=False, **kwargs)
+        self.feature_id = feature_id
+        self.feature_name = feature_name
 
     @classmethod
     def from_feature_state_model(
@@ -49,16 +50,17 @@ def from_api_flag(cls, flag_data: dict) -> "Flag":
 
 @dataclass
 class Flags:
-    flags: typing.Dict[str, BaseFlag]
+    flags: typing.Dict[str, Flag] = field(default_factory=dict)
+    default_flag_handler: typing.Callable[[str], DefaultFlag] = None
     _analytics_processor: AnalyticsProcessor = None
 
     @classmethod
     def from_feature_state_models(
         cls,
         feature_states: typing.List[FeatureStateModel],
         analytics_processor: AnalyticsProcessor,
+        default_flag_handler: typing.Callable,
         identity_id: typing.Union[str, int] = None,
-        defaults: typing.List[DefaultFlag] = None,
     ) -> "Flags":
         flags = {
             feature_state.feature.name: Flag.from_feature_state_model(
@@ -67,29 +69,31 @@ def from_feature_state_models(
             for feature_state in feature_states
         }
 
-        for default in defaults or []:
-            flags.setdefault(default.feature_name, default)
-
-        return cls(flags=flags, _analytics_processor=analytics_processor)
+        return cls(
+            flags=flags,
+            default_flag_handler=default_flag_handler,
+            _analytics_processor=analytics_processor,
+        )
 
     @classmethod
     def from_api_flags(
         cls,
         api_flags: typing.List[dict],
         analytics_processor: AnalyticsProcessor,
-        defaults: typing.List[DefaultFlag] = None,
+        default_flag_handler: typing.Callable,
     ) -> "Flags":
         flags = {
             flag_data["feature"]["name"]: Flag.from_api_flag(flag_data)
             for flag_data in api_flags
         }
 
-        for default in defaults or []:
-            flags.setdefault(default.feature_name, default)
-
-        return cls(flags=flags, _analytics_processor=analytics_processor)
+        return cls(
+            flags=flags,
+            default_flag_handler=default_flag_handler,
+            _analytics_processor=analytics_processor,
+        )
 
-    def all_flags(self) -> typing.List[BaseFlag]:
+    def all_flags(self) -> typing.List[Flag]:
         """
         Get a list of all Flag objects.
 
@@ -122,12 +126,14 @@ def get_flag(self, feature_name: str) -> BaseFlag:
         Get a specific flag given the feature name.
 
         :param feature_name: the name of the feature to retrieve the flag for.
-        :return: Flag object.
+        :return: BaseFlag object.
         :raises FlagsmithClientError: if feature doesn't exist
         """
         try:
             flag = self.flags[feature_name]
         except KeyError:
+            if self.default_flag_handler:
+                return self.default_flag_handler(feature_name)
             raise FlagsmithClientError("Feature does not exist: %s" % feature_name)
 
         if self._analytics_processor and hasattr(flag, "feature_id"):