fix: improve strict public API types

.changeset/quiet-horses-type.md

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

.github/scripts/check_strict_types.sh

@@ -0,0 +1,69 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+PYTHON_VERSION="${PYTHON_VERSION:-3.11}"
+tmp="$(mktemp -d)"
+trap 'rm -rf "$tmp"' EXIT
+
+python -m venv "$tmp/.venv"
+"$tmp/.venv/bin/python" -m pip install --quiet --upgrade pip
+"$tmp/.venv/bin/python" -m pip install --quiet . pyright
+
+cat > "$tmp/strict_posthog_types.py" <<'PY'
+# pyright: strict
+import atexit
+
+import posthog
+from posthog import FlagValue, Posthog
+
+client = Posthog("phc_test")
+atexit.register(client.shutdown)
+
+flag_value: FlagValue | None = client.get_feature_flag("flag", "user")
+all_flags: dict[str, FlagValue] | None = posthog.get_all_flags("user")
+enabled: bool | None = posthog.feature_enabled("flag", "user")
+payload: object | None = client.get_feature_flag_payload("flag", "user")
+
+_ = (flag_value, all_flags, enabled, payload)
+PY
+
+"$tmp/.venv/bin/python" - <<'PY' > "$tmp/public_api_access.py"
+import inspect
+
+import posthog
+from posthog import Posthog
+
+print("# pyright: strict")
+print("import posthog")
+print("from posthog import Posthog")
+print('client = Posthog("phc_test")')
+
+for name, obj in inspect.getmembers(Posthog):
+    if name.startswith("_"):
+        continue
+    if inspect.isfunction(obj) or inspect.ismethoddescriptor(obj):
+        print(f"client_{name} = client.{name}")
+
+for name, obj in inspect.getmembers(posthog):
+    if name.startswith("_") or name.startswith("inner_"):
+        continue
+    if inspect.isfunction(obj):
+        print(f"module_{name} = posthog.{name}")
+PY
+
+cat > "$tmp/pyrightconfig.json" <<JSON
+{
+  "typeCheckingMode": "strict",
+  "pythonVersion": "$PYTHON_VERSION",
+  "venvPath": "$tmp",
+  "venv": ".venv",
+  "reportMissingTypeStubs": "error",
+  "reportPrivateImportUsage": "error",
+  "reportUnknownArgumentType": "error",
+  "reportUnknownMemberType": "error",
+  "reportUnknownVariableType": "error"
+}
+JSON
+
+cd "$tmp"
+"$tmp/.venv/bin/python" -m pyright strict_posthog_types.py public_api_access.py

.github/workflows/ci.yml

@@ -77,6 +77,23 @@ jobs:
                   python .github/scripts/test_check_public_api.py
                   make public_api_check
 
+    strict-type-smoke:
+        name: Strict type smoke
+        runs-on: ubuntu-latest
+
+        steps:
+            - uses: actions/checkout@85e6279cec87321a52edac9c87bce653a07cf6c2
+              with:
+                  fetch-depth: 1
+
+            - name: Set up Python 3.11
+              uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55
+              with:
+                  python-version: 3.11.11
+
+            - name: Check strict downstream Pyright usage
+              run: .github/scripts/check_strict_types.sh
+
     package-build:
         name: Package build
         runs-on: ubuntu-latest

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 as 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()`` /