Improve/client: update public api and refactoring

example/example.py

@@ -4,7 +4,7 @@
 
 api_key = input("Please provide an environment api key: ")
 
-flagsmith = Flagsmith(environment_id=api_key)
+flagsmith = Flagsmith(environment_key=api_key)
 
 identifier = input("Please provide an example identity: ")
 feature_name = input("Please provide an example feature name: ")

flagsmith/analytics.py

@@ -1,4 +1,5 @@
 import json
+import typing
 from datetime import datetime
 
 from requests_futures.sessions import FuturesSession
@@ -17,22 +18,26 @@ class AnalyticsProcessor:
     the Flagsmith SDK. Docs: https://docs.flagsmith.com/advanced-use/flag-analytics.
     """
 
-    def __init__(self, environment_key: str, base_api_url: str, timeout: int = 3):
+    def __init__(
+        self, base_api_url: str, http_headers: typing.Dict[str, str], timeout: int = 3
+    ):
         """
         Initialise the AnalyticsProcessor to handle sending analytics on flag usage to
         the Flagsmith API.
 
-        :param environment_key: environment key obtained from the Flagsmith UI
         :param base_api_url: base api url to override when using self hosted version
-        :param timeout: used to tell requests to stop waiting for a response after a
+        :param http_headers: All the http headers required to communicate with the server(including x-enviroment-key)
+        :param timeout(optional): used to tell requests to stop waiting for a response after a
             given number of seconds
         """
         self.analytics_endpoint = base_api_url + ANALYTICS_ENDPOINT
-        self.environment_key = environment_key
+        self.headers = http_headers
+        # Add content type if not present
+        self.headers.update({"Content-Type": "application/json"})
+
         self._last_flushed = datetime.now()
         self.analytics_data = {}
         self.timeout = timeout
-        super().__init__()
 
     def flush(self):
         """
@@ -45,10 +50,7 @@ def flush(self):
             self.analytics_endpoint,
             data=json.dumps(self.analytics_data),
             timeout=self.timeout,
-            headers={
-                "X-Environment-Key": self.environment_key,
-                "Content-Type": "application/json",
-            },
+            headers=self.headers,
         )
 
         self.analytics_data.clear()

flagsmith/flagsmith.py

@@ -1,8 +1,10 @@
 import logging
+import typing
 
 import requests
 
-from .analytics import AnalyticsProcessor
+from flagsmith.analytics import AnalyticsProcessor
+from flagsmith.utils import generate_header_content
 
 logger = logging.getLogger(__name__)
 
@@ -14,29 +16,37 @@
 
 class Flagsmith:
     def __init__(
-        self, environment_id, api=SERVER_URL, custom_headers=None, request_timeout=None
+        self,
+        environment_key: str,
+        api: str = SERVER_URL,
+        custom_headers: typing.Dict[str, str] = None,
+        request_timeout: int = None,
     ):
         """
         Initialise Flagsmith environment.
 
-        :param environment_id: environment key obtained from the Flagsmith UI
+        :param environment_key: environment key obtained from the Flagsmith UI
         :param api: (optional) api url to override when using self hosted version
         :param custom_headers: (optional) dict which will be passed in headers for each api call
         :param request_timeout: (optional) request timeout in seconds
         """
 
-        self.environment_id = environment_id
+        self.environment_key = environment_key
         self.api = api
         self.flags_endpoint = api + FLAGS_ENDPOINT
         self.identities_endpoint = api + IDENTITY_ENDPOINT
         self.traits_endpoint = api + TRAIT_ENDPOINT
         self.custom_headers = custom_headers or {}
         self.request_timeout = request_timeout
         self._analytics_processor = AnalyticsProcessor(
-            environment_id, api, self.request_timeout
+            api,
+            generate_header_content(self.environment_key, self.custom_headers),
+            self.request_timeout,
         )
 
-    def get_flags(self, identity=None):
+    def get_flags(
+        self, identity: str = None
+    ) -> typing.Optional[typing.List[typing.Mapping]]:
         """
         Get all flags for the environment or optionally provide an identity within an environment
         to get their flags. Will return overridden identity flags where given and fill in the gaps
@@ -50,12 +60,14 @@ def get_flags(self, identity=None):
         else:
             data = self._get_flags_response()
 
-        if data:
-            return data
-        else:
+        if not data:
             logger.error("Failed to get flags for environment.")
+            return None
+        return data
 
-    def get_flags_for_user(self, identity):
+    def get_flags_for_user(
+        self, identity: str
+    ) -> typing.Optional[typing.List[typing.Mapping]]:
         """
         Get all flags for a user
 
@@ -64,7 +76,7 @@ def get_flags_for_user(self, identity):
         """
         return self.get_flags(identity=identity)
 
-    def has_feature(self, feature_name):
+    def has_feature(self, feature_name: str) -> bool:
         """
         Determine if given feature exists for an environment.
 
@@ -79,7 +91,9 @@ def has_feature(self, feature_name):
 
         return False
 
-    def feature_enabled(self, feature_name, identity=None):
+    def feature_enabled(
+        self, feature_name: str, identity: str = None
+    ) -> typing.Optional[bool]:
         """
         Get enabled state of given feature for an environment.
 
@@ -100,7 +114,9 @@ def feature_enabled(self, feature_name, identity=None):
 
         return data["enabled"]
 
-    def get_value(self, feature_name, identity=None):
+    def get_value(
+        self, feature_name: str, identity: str = None
+    ) -> typing.Union[None, int, str, bool]:
         """
         Get value of given feature for an environment.
 
@@ -119,7 +135,7 @@ def get_value(self, feature_name, identity=None):
         self._analytics_processor.track_feature(feature_id)
         return data["feature_state_value"]
 
-    def get_trait(self, trait_key, identity):
+    def get_trait(self, trait_key: str, identity: str) -> typing.Optional[str]:
         """
         Get value of given trait for the identity of an environment.
 
@@ -137,7 +153,7 @@ def get_trait(self, trait_key, identity):
             if trait.get("trait_key") == trait_key:
                 return trait.get("trait_value")
 
-    def set_trait(self, trait_key, trait_value, identity):
+    def set_trait(self, trait_key: str, trait_value: str, identity: str):
         """
         Set value of given trait for the identity of an environment. Note that this will lazily create
         a new trait if the trait_key has not been seen before for this identity
@@ -148,7 +164,7 @@ def set_trait(self, trait_key, trait_value, identity):
         """
         values = [trait_key, trait_value, identity]
         if None in values or "" in values:
-            return None
+            return
 
         payload = {
             "identity": {"identifier": identity},
@@ -159,11 +175,13 @@ def set_trait(self, trait_key, trait_value, identity):
         requests.post(
             self.traits_endpoint,
             json=payload,
-            headers=self._generate_header_content(self.custom_headers),
+            headers=generate_header_content(self.environment_key, self.custom_headers),
             timeout=self.request_timeout,
         )
 
-    def _get_flags_response(self, feature_name=None, identity=None):
+    def _get_flags_response(
+        self, feature_name: str = None, identity: str = None
+    ) -> typing.Optional[typing.Any]:
         """
         Private helper method to hit the flags endpoint
 
@@ -179,14 +197,18 @@ def _get_flags_response(self, feature_name=None, identity=None):
                 response = requests.get(
                     self.identities_endpoint,
                     params=params,
-                    headers=self._generate_header_content(self.custom_headers),
+                    headers=generate_header_content(
+                        self.environment_key, self.custom_headers
+                    ),
                     timeout=self.request_timeout,
                 )
             else:
                 response = requests.get(
                     self.flags_endpoint,
                     params=params,
-                    headers=self._generate_header_content(self.custom_headers),
+                    headers=generate_header_content(
+                        self.environment_key, self.custom_headers
+                    ),
                     timeout=self.request_timeout,
                 )
 
@@ -205,15 +227,3 @@ def _get_flags_response(self, feature_name=None, identity=None):
                 "Got error getting response from API. Error message was %s" % e
             )
             return None
-
-    def _generate_header_content(self, headers=None):
-        """
-        Generates required header content for accessing API
-
-        :param headers: (optional) dictionary of other required header values
-        :return: dictionary with required environment header appended to it
-        """
-        headers = headers or {}
-
-        headers["X-Environment-Key"] = self.environment_id
-        return headers

flagsmith/utils.py

@@ -0,0 +1,16 @@
+import typing
+
+
+def generate_header_content(
+    environment_key: str, headers: typing.Dict[str, str] = None
+) -> typing.Dict[str, str]:
+    """
+    Generates required header content for accessing API
+
+    :param headers: (optional) dictionary of other required header values
+    :return: dictionary with required environment header appended to it
+    """
+    headers = headers or {}
+
+    headers["X-Environment-Key"] = environment_key
+    return headers

tests/conftest.py

@@ -1,10 +1,11 @@
 import pytest
 
 from flagsmith.analytics import AnalyticsProcessor
+from flagsmith.utils import generate_header_content
 
 
 @pytest.fixture
 def analytics_processor():
     return AnalyticsProcessor(
-        environment_key="test_key", base_api_url="http://test_url"
+        base_api_url="http://test_url", http_headers=generate_header_content("test_key")
     )

tests/test_flagsmith.py

@@ -61,7 +61,7 @@ class FlagsmithTestCase(TestCase):
     test_environment_key = "test-env-key"
 
     def setUp(self) -> None:
-        self.bt = Flagsmith(environment_id=self.test_environment_key, api=TEST_API_URL)
+        self.bt = Flagsmith(environment_key=self.test_environment_key, api=TEST_API_URL)
 
     @mock.patch(
         "flagsmith.flagsmith.requests.get",

tests/test_utils.py

@@ -0,0 +1,23 @@
+from flagsmith.utils import generate_header_content
+
+
+def test_generate_header_content_returns_dict_with_api_key():
+    # Given
+    environment_key = "api_key"
+
+    # When
+    headers = generate_header_content(environment_key)
+
+    # Then
+    assert headers == {"X-Environment-Key": environment_key}
+
+
+def test_generate_header_content_appends_headers_to_header_arg_if_present():
+    # Given
+    environment_key = "api_key"
+    given_headers = {"Host": "test"}
+    # When
+    headers = generate_header_content(environment_key, given_headers)
+
+    # Then
+    assert headers == {"X-Environment-Key": environment_key, **given_headers}