fix: improve strict public API types

Author: marandaneto

.changeset/quiet-horses-type.md

@@ -0,0 +1,5 @@
+---
+'pypi/posthog': patch
+---
+
+Improve strict Pyright coverage for public PostHog APIs.

posthog/__init__.py

@@ -3,7 +3,7 @@
 
 from typing_extensions import Unpack
 
-from posthog.args import ExceptionArg, OptionalCaptureArgs, OptionalSetArgs
+from posthog.args import ID_TYPES, ExceptionArg, OptionalCaptureArgs, OptionalSetArgs
 from posthog.client import Client
 from posthog.exception_capture import ExceptionCapture
 from posthog.contexts import (
@@ -61,7 +61,8 @@
 )
 from posthog.types import (
     BeforeSendCallback as BeforeSendCallback,
-    FeatureFlag,
+    FeatureFlag as FeatureFlag,
+    FlagValue,
     FlagsAndPayloads,
 )
 from posthog.types import (
@@ -198,7 +199,7 @@ def set_capture_exception_code_variables_context(enabled: bool):
     return inner_set_capture_exception_code_variables_context(enabled)
 
 
-def set_code_variables_mask_patterns_context(mask_patterns: list):
+def set_code_variables_mask_patterns_context(mask_patterns: list[str]):
     """
     Override code-variable mask patterns for exceptions in the current context.
 
@@ -212,7 +213,7 @@ def set_code_variables_mask_patterns_context(mask_patterns: list):
     return inner_set_code_variables_mask_patterns_context(mask_patterns)
 
 
-def set_code_variables_ignore_patterns_context(ignore_patterns: list):
+def set_code_variables_ignore_patterns_context(ignore_patterns: list[str]):
     """
     Override code-variable ignore patterns for exceptions in the current context.
 
@@ -491,15 +492,14 @@ def set_once(**kwargs: Unpack[OptionalSetArgs]) -> Optional[str]:
 
 
 def group_identify(
-    group_type,  # type: str
-    group_key,  # type: str
-    properties=None,  # type: Optional[Dict]
-    timestamp=None,  # type: Optional[datetime.datetime]
-    uuid=None,  # type: Optional[str]
-    disable_geoip=None,  # type: Optional[bool]
-    distinct_id=None,  # type: Optional[str]
-):
-    # type: (...) -> Optional[str]
+    group_type: str,
+    group_key: str,
+    properties: Optional[Dict[str, Any]] = None,
+    timestamp: Optional[datetime.datetime] = None,
+    uuid: Optional[str] = None,
+    disable_geoip: Optional[bool] = None,
+    distinct_id: Optional[ID_TYPES] = None,
+) -> Optional[str]:
     """
     Set properties on a group.
 
@@ -538,13 +538,12 @@ def group_identify(
 
 
 def alias(
-    previous_id,  # type: str
-    distinct_id,  # type: str
-    timestamp=None,  # type: Optional[datetime.datetime]
-    uuid=None,  # type: Optional[str]
-    disable_geoip=None,  # type: Optional[bool]
-):
-    # type: (...) -> Optional[str]
+    previous_id: str,
+    distinct_id: str,
+    timestamp: Optional[datetime.datetime] = None,
+    uuid: Optional[str] = None,
+    disable_geoip: Optional[bool] = None,
+) -> Optional[str]:
     """
     Associate user behaviour before and after they e.g. register, login, or perform some other identifying action.
 
@@ -610,17 +609,16 @@ def capture_exception(
 
 
 def feature_enabled(
-    key,  # type: str
-    distinct_id,  # type: str
-    groups=None,  # type: Optional[dict]
-    person_properties=None,  # type: Optional[dict]
-    group_properties=None,  # type: Optional[dict]
-    only_evaluate_locally=False,  # type: bool
-    send_feature_flag_events=True,  # type: bool
-    disable_geoip=None,  # type: Optional[bool]
-    device_id=None,  # type: Optional[str]
-):
-    # type: (...) -> bool
+    key: str,
+    distinct_id: ID_TYPES,
+    groups: Optional[Dict[str, str]] = None,
+    person_properties: Optional[Dict[str, Any]] = None,
+    group_properties: Optional[Dict[str, Dict[str, Any]]] = None,
+    only_evaluate_locally: bool = False,
+    send_feature_flag_events: bool = True,
+    disable_geoip: Optional[bool] = None,
+    device_id: Optional[str] = None,
+) -> Optional[bool]:
     """
     Use feature flags to enable or disable features for users.
 
@@ -664,16 +662,16 @@ def feature_enabled(
 
 
 def get_feature_flag(
-    key,  # type: str
-    distinct_id,  # type: str
-    groups=None,  # type: Optional[dict]
-    person_properties=None,  # type: Optional[dict]
-    group_properties=None,  # type: Optional[dict]
-    only_evaluate_locally=False,  # type: bool
-    send_feature_flag_events=True,  # type: bool
-    disable_geoip=None,  # type: Optional[bool]
-    device_id=None,  # type: Optional[str]
-) -> Optional[FeatureFlag]:
+    key: str,
+    distinct_id: ID_TYPES,
+    groups: Optional[Dict[str, str]] = None,
+    person_properties: Optional[Dict[str, Any]] = None,
+    group_properties: Optional[Dict[str, Dict[str, Any]]] = None,
+    only_evaluate_locally: bool = False,
+    send_feature_flag_events: bool = True,
+    disable_geoip: Optional[bool] = None,
+    device_id: Optional[str] = None,
+) -> Optional[FlagValue]:
     """
     Get feature flag variant for users. Used with experiments.
 
@@ -717,15 +715,15 @@ def get_feature_flag(
 
 
 def get_all_flags(
-    distinct_id,  # type: str
-    groups=None,  # type: Optional[dict]
-    person_properties=None,  # type: Optional[dict]
-    group_properties=None,  # type: Optional[dict]
-    only_evaluate_locally=False,  # type: bool
-    disable_geoip=None,  # type: Optional[bool]
-    device_id=None,  # type: Optional[str]
-    flag_keys_to_evaluate=None,  # type: Optional[list[str]]
-) -> Optional[dict[str, FeatureFlag]]:
+    distinct_id: ID_TYPES,
+    groups: Optional[Dict[str, str]] = None,
+    person_properties: Optional[Dict[str, Any]] = None,
+    group_properties: Optional[Dict[str, Dict[str, Any]]] = None,
+    only_evaluate_locally: bool = False,
+    disable_geoip: Optional[bool] = None,
+    device_id: Optional[str] = None,
+    flag_keys_to_evaluate: Optional[list[str]] = None,
+) -> Optional[dict[str, FlagValue]]:
     """
     Get all flags for a given user.
 
@@ -765,17 +763,16 @@ def get_all_flags(
 
 
 def get_feature_flag_result(
-    key,
-    distinct_id,
-    groups=None,  # type: Optional[dict]
-    person_properties=None,  # type: Optional[dict]
-    group_properties=None,  # type: Optional[dict]
-    only_evaluate_locally=False,
-    send_feature_flag_events=True,
-    disable_geoip=None,  # type: Optional[bool]
-    device_id=None,  # type: Optional[str]
-):
-    # type: (...) -> Optional[FeatureFlagResult]
+    key: str,
+    distinct_id: ID_TYPES,
+    groups: Optional[Dict[str, str]] = None,
+    person_properties: Optional[Dict[str, Any]] = None,
+    group_properties: Optional[Dict[str, Dict[str, Any]]] = None,
+    only_evaluate_locally: bool = False,
+    send_feature_flag_events: bool = True,
+    disable_geoip: Optional[bool] = None,
+    device_id: Optional[str] = None,
+) -> Optional[FeatureFlagResult]:
     """
     Get a FeatureFlagResult object which contains the flag result and payload.
 
@@ -821,17 +818,17 @@ def get_feature_flag_result(
 
 
 def get_feature_flag_payload(
-    key,
-    distinct_id,
-    match_value=None,
-    groups=None,  # type: Optional[dict]
-    person_properties=None,  # type: Optional[dict]
-    group_properties=None,  # type: Optional[dict]
-    only_evaluate_locally=False,
-    send_feature_flag_events=True,
-    disable_geoip=None,  # type: Optional[bool]
-    device_id=None,  # type: Optional[str]
-) -> Optional[str]:
+    key: str,
+    distinct_id: ID_TYPES,
+    match_value: Optional[FlagValue] = None,
+    groups: Optional[Dict[str, str]] = None,
+    person_properties: Optional[Dict[str, Any]] = None,
+    group_properties: Optional[Dict[str, Dict[str, Any]]] = None,
+    only_evaluate_locally: bool = False,
+    send_feature_flag_events: bool = True,
+    disable_geoip: Optional[bool] = None,
+    device_id: Optional[str] = None,
+) -> Optional[Any]:
     """
     Get the payload associated with a feature flag value.
 
@@ -869,7 +866,7 @@ def get_feature_flag_payload(
 
 
 def get_remote_config_payload(
-    key,  # type: str
+    key: str,
 ):
     """Get the payload for a remote config feature flag.
 
@@ -889,14 +886,14 @@ def get_remote_config_payload(
 
 
 def get_all_flags_and_payloads(
-    distinct_id,
-    groups=None,  # type: Optional[dict]
-    person_properties=None,  # type: Optional[dict]
-    group_properties=None,  # type: Optional[dict]
-    only_evaluate_locally=False,
-    disable_geoip=None,  # type: Optional[bool]
-    device_id=None,  # type: Optional[str]
-    flag_keys_to_evaluate=None,  # type: Optional[list[str]]
+    distinct_id: ID_TYPES,
+    groups: Optional[Dict[str, str]] = None,
+    person_properties: Optional[Dict[str, Any]] = None,
+    group_properties: Optional[Dict[str, Dict[str, Any]]] = None,
+    only_evaluate_locally: bool = False,
+    disable_geoip: Optional[bool] = None,
+    device_id: Optional[str] = None,
+    flag_keys_to_evaluate: Optional[list[str]] = None,
 ) -> FlagsAndPayloads:
     """
     Get all feature flag values and payloads for a user.
@@ -932,14 +929,14 @@ def get_all_flags_and_payloads(
 
 
 def evaluate_flags(
-    distinct_id=None,  # type: Optional[str]
-    groups=None,  # type: Optional[Dict[str, str]]
-    person_properties=None,  # type: Optional[Dict[str, Any]]
-    group_properties=None,  # type: Optional[Dict[str, Dict[str, Any]]]
-    only_evaluate_locally=False,  # type: bool
-    disable_geoip=None,  # type: Optional[bool]
-    flag_keys=None,  # type: Optional[list]
-    device_id=None,  # type: Optional[str]
+    distinct_id: Optional[str] = None,
+    groups: Optional[Dict[str, str]] = None,
+    person_properties: Optional[Dict[str, Any]] = None,
+    group_properties: Optional[Dict[str, Dict[str, Any]]] = None,
+    only_evaluate_locally: bool = False,
+    disable_geoip: Optional[bool] = None,
+    flag_keys: Optional[list[str]] = None,
+    device_id: Optional[str] = None,
 ) -> FeatureFlagEvaluations:
     """Evaluate all feature flags for a user in a single call and return a
     :class:`FeatureFlagEvaluations` snapshot. Branch on ``.is_enabled()`` /