Unwrap and be happy

Author: emyller

api/features/views.py

@@ -62,9 +62,9 @@
     NestedEnvironmentPermissions,
 )
 from features.value_types import BOOLEAN, INTEGER, STRING
+from integrations.flagsmith.client import get_client
 from projects.code_references.services import (
     annotate_feature_queryset_with_code_references_summary,
-    is__code_references_ui_stats__enabled,
 )
 from projects.models import Project
 from users.models import FFAdminUser, UserPermissionGroup
@@ -221,7 +221,13 @@ def get_queryset(self):  # type: ignore[no-untyped-def]
         query_data = query_serializer.validated_data
 
         # TODO: Delete this after https://github.com/flagsmith/flagsmith/issues/6832 is resolved
-        if is__code_references_ui_stats__enabled(project.organisation):
+        organisation = project.organisation
+        flagsmith_client = get_client("local", local_eval=True)
+        flags = flagsmith_client.get_identity_flags(
+            organisation.flagsmith_identifier,
+            traits=organisation.flagsmith_on_flagsmith_api_traits,
+        )
+        if flags.is_feature_enabled("code_references_ui_stats"):
             queryset = annotate_feature_queryset_with_code_references_summary(queryset)
         else:
             queryset = queryset.annotate(

api/projects/code_references/services.py

@@ -1,4 +1,3 @@
-import typing
 from datetime import timedelta
 from urllib.parse import urljoin
 
@@ -16,7 +15,6 @@
 from django.utils import timezone
 
 from features.models import Feature
-from integrations.flagsmith.client import get_client
 from projects.code_references.constants import (
     FEATURE_FLAG_CODE_REFERENCES_RETENTION_DAYS,
 )
@@ -27,22 +25,6 @@
     VCSProvider,
 )
 
-if typing.TYPE_CHECKING:
-    from organisations.models import Organisation
-
-
-def is__code_references_ui_stats__enabled(organisation: "Organisation") -> bool:
-    """Check if code reference UI stats are enabled for an organisation.
-
-    TODO: Delete this after https://github.com/flagsmith/flagsmith/issues/6832 is resolved.
-    """
-    flagsmith_client = get_client("local", local_eval=True)
-    flags = flagsmith_client.get_identity_flags(
-        organisation.flagsmith_identifier,
-        traits=organisation.flagsmith_on_flagsmith_api_traits,
-    )
-    return flags.is_feature_enabled("code_references_ui_stats")
-
 
 def annotate_feature_queryset_with_code_references_summary(
     queryset: QuerySet[Feature],