Python: Fix Python pyright package scoping and typing remediation (#4426)

* Fix Python pyright package scoping and typing remediation

Implements issue #4407 by removing the root pyright include, adding package-level pyright includes, and resolving pyright/mypy typing issues across Python packages. Also cleans unnecessary casts and applies line-level, rule-specific ignores where external libraries are too dynamic.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Reduce pyright cost in handoff cloning

Simplify cloned_options construction in HandoffAgentExecutor to avoid expensive TypedDict narrowing/inference in _handoff.py, which was causing pyright to spend a long time in orchestrations.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix types

* Fix lint and type-check regressions

Resolve current Python package check failures across lint, pyright, and mypy after recent code changes, including purview/declarative pyright issues and multiple ruff simplification findings.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fixed hooks

* Stabilize package tests and test tasks

Resolve cross-package non-integration test failures, simplify streaming type flow, harden locale/culture handling, and standardize package test poe tasks to exclude integration tests where applicable.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* lots of small fixes

* Fix current Python test regressions

Address current failing unit tests in azure-ai, bedrock, and azure-cosmos while keeping Bedrock parsing logic inline (no new static helper methods).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* small fixes

* small fixes

* removed pydantic from json

* final updates

* fix core

* fix tests

* fix obser

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

python/packages/purview/agent_framework_purview/_client.py

@@ -6,7 +6,7 @@ import base64
 import inspect
 import json
 import logging
-from typing import Any, cast
+from typing import Any, Literal, TypeVar, overload
 from uuid import uuid4
 
 import httpx
@@ -36,6 +36,8 @@ from ._settings import PurviewSettings, get_purview_scopes
 
 logger = logging.getLogger("agent_framework.purview")
 
+ResponseT = TypeVar("ResponseT")
+
 
 class PurviewClient:
     """Async client for calling Graph Purview endpoints.
@@ -98,7 +100,7 @@ class PurviewClient:
         with get_tracer().start_as_current_span("purview.process_content"):
             token = await self._get_token(tenant_id=request.tenant_id)
             url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/processContent"
-            headers = {}
+            headers: dict[str, str] = {}
             # Add If-None-Match header if scope_identifier is present
             if hasattr(request, "scope_identifier") and request.scope_identifier:
                 headers["If-None-Match"] = request.scope_identifier
@@ -106,21 +108,23 @@ class PurviewClient:
             if hasattr(request, "process_inline") and request.process_inline:
                 headers["Prefer"] = "evaluateInline"
 
-            response = await self._post(
+            response: ProcessContentResponse | tuple[ProcessContentResponse, httpx.Headers] = await self._post(
                 url, request, ProcessContentResponse, token, headers=headers, return_response=True
             )
 
             if isinstance(response, tuple) and len(response) == 2:
                 response_obj, _ = response
-                return cast(ProcessContentResponse, response_obj)
+                return response_obj
 
-            return cast(ProcessContentResponse, response)
+            return response
 
     async def get_protection_scopes(self, request: ProtectionScopesRequest) -> ProtectionScopesResponse:
         with get_tracer().start_as_current_span("purview.get_protection_scopes"):
             token = await self._get_token()
             url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/protectionScopes/compute"
-            response = await self._post(url, request, ProtectionScopesResponse, token, return_response=True)
+            response: ProtectionScopesResponse | tuple[ProtectionScopesResponse, httpx.Headers] = await self._post(
+                url, request, ProtectionScopesResponse, token, return_response=True
+            )
 
             # Extract etag from response headers
             if isinstance(response, tuple) and len(response) == 2:
@@ -128,25 +132,47 @@ class PurviewClient:
                 if "etag" in headers:
                     etag_value = headers["etag"].strip('"')
                     response_obj.scope_identifier = etag_value
-                return cast(ProtectionScopesResponse, response_obj)
+                return response_obj
 
-            return cast(ProtectionScopesResponse, response)
+            return response
 
     async def send_content_activities(self, request: ContentActivitiesRequest) -> ContentActivitiesResponse:
         with get_tracer().start_as_current_span("purview.send_content_activities"):
             token = await self._get_token()
             url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/activities/contentActivities"
-            return cast(ContentActivitiesResponse, await self._post(url, request, ContentActivitiesResponse, token))
+            return await self._post(url, request, ContentActivitiesResponse, token)
+
+    @overload
+    async def _post(
+        self,
+        url: str,
+        model: Any,
+        response_type: type[ResponseT],
+        token: str,
+        headers: dict[str, str] | None = None,
+        return_response: Literal[False] = False,
+    ) -> ResponseT: ...
+
+    @overload
+    async def _post(
+        self,
+        url: str,
+        model: Any,
+        response_type: type[ResponseT],
+        token: str,
+        headers: dict[str, str] | None = None,
+        return_response: Literal[True] = True,
+    ) -> tuple[ResponseT, httpx.Headers]: ...
 
     async def _post(
         self,
         url: str,
         model: Any,
-        response_type: type[Any],
+        response_type: type[ResponseT],
         token: str,
         headers: dict[str, str] | None = None,
         return_response: bool = False,
-    ) -> Any:
+    ) -> ResponseT | tuple[ResponseT, httpx.Headers]:
         if hasattr(model, "correlation_id") and not model.correlation_id:
             model.correlation_id = str(uuid4())
 
@@ -174,7 +200,7 @@ class PurviewClient:
             raise PurviewAuthenticationError(f"Auth failure {resp.status_code}: {resp.text}")
         if resp.status_code == 402:
             if self._settings.get("ignore_payment_required", False):
-                return response_type()  # type: ignore[call-arg, no-any-return]
+                return response_type()  # type: ignore[call-arg]
             raise PurviewPaymentRequiredError(f"Payment required {resp.status_code}: {resp.text}")
         if resp.status_code == 429:
             raise PurviewRateLimitError(f"Rate limited {resp.status_code}: {resp.text}")
@@ -187,18 +213,18 @@ class PurviewClient:
 
         try:
             # Prefer pydantic-style model_validate if present, else fall back to constructor.
-            if hasattr(response_type, "model_validate"):
-                response_obj = response_type.model_validate(data)  # type: ignore[no-any-return]
-            else:
-                response_obj = response_type(**data)  # type: ignore[call-arg, no-any-return]
+            model_validate = getattr(response_type, "model_validate", None)
+            response_obj = model_validate(data) if callable(model_validate) else response_type(**data)  # type: ignore[call-arg]
 
             # Extract correlation_id from response headers if response object supports it
             if "client-request-id" in resp.headers and hasattr(response_obj, "correlation_id"):
-                response_obj.correlation_id = resp.headers["client-request-id"]
-                logger.info(f"Purview response from {url} with correlation_id: {response_obj.correlation_id}")
+                response_correlation_id = resp.headers["client-request-id"]
+                response_obj.correlation_id = response_correlation_id  # pyright: ignore[reportAttributeAccessIssue]
+                logger.info(f"Purview response from {url} with correlation_id: {response_correlation_id}")
 
+            typed_response_obj = response_obj if isinstance(response_obj, response_type) else response_type(**data)
             if return_response:
-                return (response_obj, resp.headers)
-            return response_obj
+                return (typed_response_obj, resp.headers)
+            return typed_response_obj
         except Exception as ex:
             raise PurviewServiceError(f"Failed to deserialize Purview response: {ex}") from ex

python/packages/purview/agent_framework_purview/_middleware.py

@@ -67,6 +67,7 @@ class PurviewPolicyMiddleware(AgentMiddleware):
         call_next: Callable[[], Awaitable[None]],
     ) -> None:  # type: ignore[override]
         resolved_user_id: str | None = None
+        session_id: str | None = None
         try:
             # Pre (prompt) check
             session_id = self._get_agent_session_id(context)
@@ -107,7 +108,7 @@ class PurviewPolicyMiddleware(AgentMiddleware):
                 should_block_response, _ = await self._processor.process_messages(
                     context.result.messages,  # type: ignore[union-attr]
                     Activity.DOWNLOAD_TEXT,
-                    session_id=session_id,
+                    session_id=session_id_response,
                     user_id=resolved_user_id,
                 )
                 if should_block_response:
@@ -173,6 +174,7 @@ class PurviewChatPolicyMiddleware(ChatMiddleware):
         call_next: Callable[[], Awaitable[None]],
     ) -> None:  # type: ignore[override]
         resolved_user_id: str | None = None
+        session_id: str | None = None
         try:
             session_id = context.options.get("conversation_id") if context.options else None
             should_block_prompt, resolved_user_id = await self._processor.process_messages(

python/packages/purview/agent_framework_purview/_models.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Mapping, MutableMapping, Sequence
+from collections.abc import Iterable, Mapping, MutableMapping, Sequence
 from datetime import datetime
 from enum import Enum, Flag, auto
 from typing import Any, ClassVar, TypeVar, cast
@@ -60,6 +60,23 @@ _PROTECTION_SCOPE_ACTIVITIES_SERIALIZE_ORDER: list[tuple[str, ProtectionScopeAct
 ]
 
 
+def _as_object_list(value: object) -> list[object] | None:
+    if not isinstance(value, (list, tuple, set)):
+        return None
+    return list(cast(Iterable[object], value))
+
+
+def _as_str_dict(value: object) -> dict[str, str]:
+    if not isinstance(value, dict):
+        return {}
+
+    aliases: dict[str, str] = {}
+    for raw_key, raw_value in cast(dict[object, object], value).items():
+        if isinstance(raw_key, str) and isinstance(raw_value, str):
+            aliases[raw_key] = raw_value
+    return aliases
+
+
 def deserialize_flag(
     value: object, mapping: Mapping[str, FlagT], enum_cls: type[FlagT]
 ) -> FlagT | None:  # pragma: no cover
@@ -82,8 +99,11 @@ def deserialize_flag(
         if not raw:
             return enum_cls(0)
         parts.extend([p.strip() for p in raw.split(",") if p.strip()])
-    elif isinstance(value, (list, tuple, set)):
-        for item in value:
+    else:
+        iterable_items = _as_object_list(value)
+        if iterable_items is None:
+            return None
+        for item in iterable_items:
             if isinstance(item, str):
                 parts.extend([p.strip() for p in item.split(",") if p.strip()])
             elif isinstance(item, enum_cls):
@@ -93,8 +113,6 @@ def deserialize_flag(
                     flag_value |= enum_cls(item)
                 except Exception:
                     logger.warning(f"Failed to convert int {item} to {enum_cls.__name__}")
-    else:
-        return None
 
     for part in parts:
         member = mapping.get(part)
@@ -196,10 +214,10 @@ class _AliasSerializable(SerializationMixin):
         # Collect all aliases from parent classes too
         all_aliases: dict[str, str] = {}
         for cls in type(self).__mro__:
-            if hasattr(cls, "_ALIASES") and isinstance(cls._ALIASES, dict):
-                for internal, external in cls._ALIASES.items():
-                    if external not in all_aliases:
-                        all_aliases[external] = internal
+            aliases_obj = _as_str_dict(getattr(cls, "_ALIASES", None))
+            for internal, external in aliases_obj.items():
+                if external not in all_aliases:
+                    all_aliases[external] = internal
 
         # Normalize all aliased keys in kwargs
         for external, internal in all_aliases.items():
@@ -248,11 +266,11 @@ class _AliasSerializable(SerializationMixin):
         # Collect all aliases from class hierarchy
         all_aliases: dict[str, str] = {}
         for cls in type(self).__mro__:
-            if hasattr(cls, "_ALIASES") and isinstance(cls._ALIASES, dict):
-                # Parent aliases first (will be overridden by child if same key)
-                for internal, external in cls._ALIASES.items():
-                    if internal not in all_aliases:
-                        all_aliases[internal] = external
+            aliases_obj = _as_str_dict(getattr(cls, "_ALIASES", None))
+            # Parent aliases first (will be overridden by child if same key)
+            for internal, external in aliases_obj.items():
+                if internal not in all_aliases:
+                    all_aliases[internal] = external
 
         if not all_aliases:
             return base
@@ -836,17 +854,15 @@ class ProcessContentResponse(_AliasSerializable):
         # Convert to objects
         converted_policy_actions: list[DlpActionInfo] | None = None
         if policy_actions is not None:
-            converted_policy_actions = cast(
-                list[DlpActionInfo],
-                [p if isinstance(p, DlpActionInfo) else DlpActionInfo(**p) for p in policy_actions],
-            )
+            converted_policy_actions = [
+                p if isinstance(p, DlpActionInfo) else DlpActionInfo(**p) for p in policy_actions
+            ]
 
         converted_processing_errors: list[ProcessingError] | None = None
         if processing_errors is not None:
-            converted_processing_errors = cast(
-                list[ProcessingError],
-                [pe if isinstance(pe, ProcessingError) else ProcessingError(**pe) for pe in processing_errors],
-            )
+            converted_processing_errors = [
+                pe if isinstance(pe, ProcessingError) else ProcessingError(**pe) for pe in processing_errors
+            ]
 
         super().__init__(**kwargs)
         self.id = id
@@ -885,17 +901,15 @@ class PolicyScope(_AliasSerializable):
         # Convert nested objects
         converted_locations: list[PolicyLocation] | None = None
         if locations is not None:
-            converted_locations = cast(
-                list[PolicyLocation],
-                [loc if isinstance(loc, PolicyLocation) else PolicyLocation(**loc) for loc in locations],
-            )
+            converted_locations = [
+                loc if isinstance(loc, PolicyLocation) else PolicyLocation(**loc) for loc in locations
+            ]
 
         converted_policy_actions: list[DlpActionInfo] | None = None
         if policy_actions is not None:
-            converted_policy_actions = cast(
-                list[DlpActionInfo],
-                [p if isinstance(p, DlpActionInfo) else DlpActionInfo(**p) for p in policy_actions],
-            )
+            converted_policy_actions = [
+                p if isinstance(p, DlpActionInfo) else DlpActionInfo(**p) for p in policy_actions
+            ]
 
         # Call parent without explicit params with aliases
         super().__init__(**kwargs)
@@ -947,9 +961,7 @@ class ProtectionScopesResponse(_AliasSerializable):
 
         converted_scopes: list[PolicyScope] | None = None
         if scopes is not None:
-            converted_scopes = cast(
-                list[PolicyScope], [s if isinstance(s, PolicyScope) else PolicyScope(**s) for s in scopes]
-            )
+            converted_scopes = [s if isinstance(s, PolicyScope) else PolicyScope(**s) for s in scopes]
 
         # Don't pass parameters that have aliases - let parent normalize them
         super().__init__(**kwargs)

python/packages/purview/agent_framework_purview/_processor.py

@@ -177,14 +177,13 @@ class ScopedContentProcessor:
             else:
                 raise ValueError("App location not provided or inferable")
 
+            app_name = self._settings.get("app_name") or "Unknown"
             protected_app = ProtectedAppMetadata(
-                name=self._settings["app_name"],
+                name=app_name,
                 version=self._settings.get("app_version", "Unknown"),
                 application_location=policy_location,
             )
-            integrated_app = IntegratedAppMetadata(
-                name=self._settings["app_name"], version=self._settings.get("app_version", "Unknown")
-            )
+            integrated_app = IntegratedAppMetadata(name=app_name, version=self._settings.get("app_version", "Unknown"))
             device_meta = DeviceMetadata(
                 operating_system_specifications=OperatingSystemSpecifications(
                     operating_system_platform="Unknown", operating_system_version="Unknown"
@@ -234,9 +233,9 @@ class ScopedContentProcessor:
         if cached_ps_resp is not None and isinstance(cached_ps_resp, ProtectionScopesResponse):
             ps_resp = cached_ps_resp
         else:
+            ttl = self._settings.get("cache_ttl_seconds")
+            ttl_seconds = ttl if ttl is not None else 14400
             try:
-                ttl = self._settings.get("cache_ttl_seconds")
-                ttl_seconds = ttl if ttl is not None else 14400
                 ps_resp = await self._client.get_protection_scopes(ps_req)
                 await self._cache.set(cache_key, ps_resp, ttl_seconds=ttl_seconds)
             except PurviewPaymentRequiredError as ex:

python/packages/purview/pyproject.toml

@@ -60,6 +60,7 @@ omit = [
 
 [tool.pyright]
 extends = "../../pyproject.toml"
+include = ["agent_framework_purview"]
 
 [tool.mypy]
 plugins = ['pydantic.mypy']
@@ -85,7 +86,7 @@ include = "../../shared_tasks.toml"
 
 [tool.poe.tasks]
 mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_purview"
-test = "pytest --cov=agent_framework_purview --cov-report=term-missing:skip-covered tests"
+test = "pytest -m \"not integration\" --cov=agent_framework_purview --cov-report=term-missing:skip-covered tests"
 
 [build-system]
 requires = ["flit-core >= 3.9,<4.0"]