ditch pydantic dependency

Author: khvn26

flagsmith/api/__init__.py

@@ -207,7 +207,7 @@ def handle_stream_event(self, event: StreamEvent) -> None:
             raise ValueError(
                 "Cannot handle stream events before retrieving initial environment"
             )
-        if event.updated_at > environment_updated_at:
+        if event["updated_at"] > environment_updated_at:
             self.update_environment()
 
     def get_environment_flags(self) -> Flags:

flagsmith/flagsmith.py

@@ -1,15 +1,19 @@
+import json
 import typing
 from collections import defaultdict
 from datetime import datetime, timezone
 from operator import itemgetter
 
+import sseclient
 from flag_engine.context.types import (
     EvaluationContext,
     FeatureContext,
     SegmentContext,
     SegmentRule,
 )
 
+from flagsmith.types import StreamEvent
+
 OverrideKey = typing.Tuple[
     str,
     str,
@@ -19,13 +23,24 @@
 OverridesKey = typing.Tuple[OverrideKey, ...]
 
 
+def map_sse_event_to_stream_event(event: sseclient.Event) -> StreamEvent:
+    event_data = json.loads(event.data)
+    return {
+        "updated_at": datetime.fromtimestamp(
+            event_data["updated_at"],
+            tz=timezone.utc,
+        )
+    }
+
+
 def map_environment_document_to_environment_updated_at(
     environment_document: dict[str, typing.Any],
 ) -> datetime:
-    updated_at = datetime.fromisoformat(environment_document["updated_at"])
-    if updated_at.tzinfo is None:
-        return updated_at.astimezone(tz=timezone.utc)
-    return updated_at
+    if (
+        updated_at := datetime.fromisoformat(environment_document["updated_at"])
+    ).tzinfo is None:
+        return updated_at.replace(tzinfo=timezone.utc)
+    return updated_at.astimezone(tz=timezone.utc)
 
 
 def map_environment_document_to_context(

flagsmith/mappers.py

@@ -3,15 +3,13 @@
 import typing
 from typing import Callable, Optional
 
-import pydantic
 import requests
 import sseclient
 
-logger = logging.getLogger(__name__)
-
+from flagsmith.mappers import map_sse_event_to_stream_event
+from flagsmith.types import StreamEvent
 
-class StreamEvent(pydantic.BaseModel):
-    updated_at: pydantic.AwareDatetime
+logger = logging.getLogger(__name__)
 
 
 class EventStreamManager(threading.Thread):
@@ -40,9 +38,9 @@ def run(self) -> None:
                 ) as response:
                     sse_client = sseclient.SSEClient(chunk for chunk in response)
                     for event in sse_client.events():
-                        self.on_event(StreamEvent.model_validate_json(event.data))
+                        self.on_event(map_sse_event_to_stream_event(event))
 
-            except (requests.RequestException, pydantic.ValidationError):
+            except (requests.RequestException, ValueError, TypeError):
                 logger.exception("Error opening or reading from the event stream")
 
     def stop(self) -> None:

flagsmith/streaming_manager.py

@@ -1,4 +1,5 @@
 import typing
+from datetime import datetime
 
 from flag_engine.engine import ContextValue
 from typing_extensions import NotRequired, TypeAlias
@@ -17,6 +18,10 @@
 ]
 
 
+class StreamEvent(typing.TypedDict):
+    updated_at: datetime
+
+
 class TraitConfig(typing.TypedDict):
     value: ContextValue
     transient: bool

flagsmith/types.py

@@ -15,7 +15,6 @@ requests = "^2.32.3"
 requests-futures = "^1.0.1"
 flagsmith-flag-engine = { git = "https://github.com/Flagsmith/flagsmith-engine/", rev = "9ab732a20975fefa5996d89a2cbfd95ffb52c3d3" }
 sseclient-py = "^1.8.0"
-pydantic = "^2"
 
 [tool.poetry.group.dev]
 optional = true
@@ -31,7 +30,6 @@ types-requests = "^2.32"
 pyfakefs = "^5.9.2"
 
 [tool.mypy]
-plugins = ["pydantic.mypy"]
 exclude = ["example/*"]
 
 [tool.black]

poetry.lock

@@ -37,25 +37,23 @@ def test_stream_manager_handles_timeout(
 def test_environment_updates_on_recent_event(
     server_api_key: str, mocker: MockerFixture
 ) -> None:
-    stream_updated_at = datetime(2020, 1, 1, 1, 1, 2)
+    stream_updated_at = datetime(2020, 1, 1, 1, 1, 2, tzinfo=timezone.utc)
     environment_updated_at = datetime(2020, 1, 1, 1, 1, 1, tzinfo=timezone.utc)
 
     mocker.patch("flagsmith.Flagsmith.update_environment")
 
     flagsmith = Flagsmith(environment_key=server_api_key)
     flagsmith._evaluation_context = MagicMock()
     flagsmith._environment_updated_at = environment_updated_at
-    flagsmith.handle_stream_event(
-        event=StreamEvent(updated_at=stream_updated_at.timestamp())
-    )
+    flagsmith.handle_stream_event(event=StreamEvent(updated_at=stream_updated_at))
     assert isinstance(flagsmith.update_environment, Mock)
     flagsmith.update_environment.assert_called_once()
 
 
 def test_environment_does_not_update_on_past_event(
     server_api_key: str, mocker: MockerFixture
 ) -> None:
-    stream_updated_at = datetime(2020, 1, 1, 1, 1, 1)
+    stream_updated_at = datetime(2020, 1, 1, 1, 1, 1, tzinfo=timezone.utc)
     environment_updated_at = datetime(2020, 1, 1, 1, 1, 2, tzinfo=timezone.utc)
 
     mocker.patch("flagsmith.Flagsmith.update_environment")
@@ -64,17 +62,15 @@ def test_environment_does_not_update_on_past_event(
     flagsmith._evaluation_context = MagicMock()
     flagsmith._environment_updated_at = environment_updated_at
 
-    flagsmith.handle_stream_event(
-        event=StreamEvent(updated_at=stream_updated_at.timestamp())
-    )
+    flagsmith.handle_stream_event(event=StreamEvent(updated_at=stream_updated_at))
     assert isinstance(flagsmith.update_environment, Mock)
     flagsmith.update_environment.assert_not_called()
 
 
 def test_environment_does_not_update_on_same_event(
     server_api_key: str, mocker: MockerFixture
 ) -> None:
-    stream_updated_at = datetime(2020, 1, 1, 1, 1, 1)
+    stream_updated_at = datetime(2020, 1, 1, 1, 1, 1, tzinfo=timezone.utc)
     environment_updated_at = datetime(2020, 1, 1, 1, 1, 1, tzinfo=timezone.utc)
 
     mocker.patch("flagsmith.Flagsmith.update_environment")
@@ -83,8 +79,6 @@ def test_environment_does_not_update_on_same_event(
     flagsmith._evaluation_context = MagicMock()
     flagsmith._environment_updated_at = environment_updated_at
 
-    flagsmith.handle_stream_event(
-        event=StreamEvent(updated_at=stream_updated_at.timestamp())
-    )
+    flagsmith.handle_stream_event(event=StreamEvent(updated_at=stream_updated_at))
     assert isinstance(flagsmith.update_environment, Mock)
     flagsmith.update_environment.assert_not_called()