From 4b64b6728b072409052bf1ad66eab682d3e3265d Mon Sep 17 00:00:00 2001 From: mameikagou Date: Sat, 22 Aug 2026 10:59:24 +0800 Subject: [PATCH 1/2] fix: resolve top-level resource types for mypy --- py/src/braintrust/__init__.py | 120 +++++++++++++++++- .../type_tests/test_public_exports.py | 22 +++- 2 files changed, 135 insertions(+), 7 deletions(-) diff --git a/py/src/braintrust/__init__.py b/py/src/braintrust/__init__.py index 1e66fced..5e0e5990 100644 --- a/py/src/braintrust/__init__.py +++ b/py/src/braintrust/__init__.py @@ -67,7 +67,125 @@ def is_equal(expected, output): from .framework2 import * from .functions.invoke import * from .functions.stream import * -from .generated_types import * +from .generated_types import ( # noqa: F401 + Acl, + AclObjectType, + Agent, + AISecret, + AnyModelParams, + ApiKey, + AsyncScoringControl, + AsyncScoringState, + AttachmentReference, + AttachmentStatus, + AutomationStatus, + BatchedFacetData, + BraintrustAttachmentReference, + BraintrustModelParams, + CallEvent, + ChatCompletionContentPart, + ChatCompletionContentPartFileFile, + ChatCompletionContentPartFileWithTitle, + ChatCompletionContentPartImageWithTitle, + ChatCompletionContentPartText, + ChatCompletionContentPartTextWithTitle, + ChatCompletionMessageParam, + ChatCompletionMessageReasoning, + ChatCompletionMessageToolCall, + ChatCompletionOpenAIMessageParam, + ChatCompletionTool, + CodeBundle, + DatasetEvent, + DatasetSnapshot, + EnvVar, + ExperimentEvent, + ExtendedSavedFunctionId, + ExternalAttachmentReference, + FacetData, + FacetPreprocessorId, + Function, + FunctionData, + FunctionFormat, + FunctionId, + FunctionIdRef, + FunctionObjectType, + FunctionOutputType, + FunctionTypeEnum, + FunctionTypeEnumNullish, + GitMetadataSettings, + GraphData, + GraphEdge, + GraphNode, + Group, + GroupScope, + IfExists, + ImageRenderingMode, + InvokeFunction, + InvokeParent, + MCPServer, + MessageRole, + ModelParams, + NullableSavedFunctionId, + ObjectReference, + ObjectReferenceNullish, + OnlineScoreConfig, + Organization, + OrgAutomation, + Permission, + PreprocessorId, + ProjectAutomation, + ProjectGroup, + ProjectLogsEvent, + ProjectScore, + ProjectScoreCategories, + ProjectScoreCategory, + ProjectScoreCondition, + ProjectScoreConfig, + ProjectScoreType, + ProjectSettings, + ProjectTag, + PromptBlockData, + PromptBlockDataNullish, + PromptData, + PromptDataNullish, + PromptOptions, + PromptOptionsNullish, + PromptParserNullish, + PromptSessionEvent, + RepoInfo, + ResponseFormat, + ResponseFormatJsonSchema, + ResponseFormatNullish, + RetentionObjectType, + Role, + RunEval, + SavedFunctionId, + ServiceToken, + SpanAttributes, + SpanIFrame, + SpanScope, + SpanType, + SSEConsoleEventData, + SSEProgressEventData, + StreamingMode, + ToolFunctionDefinition, + TopicAutomationConfig, + TopicAutomationDataScope, + TopicAutomationFacetModel, + TopicDigestAutomationConfig, + TopicMapData, + TopicMapFunctionAutomation, + TopicMapGenerationSettings, + TraceScope, + TriggeredFunctionState, + UploadStatus, + User, + View, + ViewData, + ViewDataSearch, + ViewOptions, + WindowedAutomationConfig, +) from .integrations.ai_sdk import setup_ai_sdk as setup_ai_sdk from .integrations.anthropic import wrap_anthropic as wrap_anthropic from .integrations.instructor import wrap_instructor as wrap_instructor diff --git a/py/src/braintrust/type_tests/test_public_exports.py b/py/src/braintrust/type_tests/test_public_exports.py index 928f6716..66443861 100644 --- a/py/src/braintrust/type_tests/test_public_exports.py +++ b/py/src/braintrust/type_tests/test_public_exports.py @@ -1,10 +1,8 @@ -"""Regression test for pyright's ``reportPrivateImportUsage`` on top-level ``braintrust`` symbols. +"""Regression tests for top-level ``braintrust`` symbols. -Without PEP 484 ``as``-aliasing in ``braintrust/__init__.py``, pyright flags -``from braintrust import auto_instrument`` (and peers) as private in a -``py.typed`` consumer. The local ``pyrightconfig.json`` turns the rule into -an error so this file breaks ``nox -s test_types`` if someone regresses the -aliasing pattern. +The static resource check keeps mypy from resolving generated ``TypedDict`` +names instead of the public runtime classes. The runtime checks cover PEP 484 +aliasing for pyright's ``reportPrivateImportUsage`` rule. """ import braintrust @@ -31,6 +29,18 @@ ] +def accepts_public_resource_types( + experiment: braintrust.Experiment, + dataset: braintrust.Dataset, + project: braintrust.Project, + prompt: braintrust.Prompt, +) -> None: + experiment.fetch() + dataset.fetch() + _ = project.name + prompt.build() + + @pytest.mark.parametrize("name,imported", _PUBLIC_SYMBOLS) def test_top_level_public_symbol(name: str, imported: object) -> None: assert callable(imported) From 752195b7c2f42e024f27d1776f8045a969cd4779 Mon Sep 17 00:00:00 2001 From: mameikagou Date: Mon, 24 Aug 2026 11:35:14 +0800 Subject: [PATCH 2/2] fix: keep generated root exports synchronized --- py/src/braintrust/__init__.py | 144 +++--------------- .../type_tests/test_public_exports.py | 30 +++- 2 files changed, 53 insertions(+), 121 deletions(-) diff --git a/py/src/braintrust/__init__.py b/py/src/braintrust/__init__.py index 5e0e5990..f568fdb0 100644 --- a/py/src/braintrust/__init__.py +++ b/py/src/braintrust/__init__.py @@ -50,6 +50,17 @@ def is_equal(expected, output): # Check env var at import time for auto-instrumentation import os +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + # These names must precede the generated wildcard so type checkers keep the + # runtime resource classes for the four names shared by both modules. + from .logger import Dataset as Dataset # noqa: I001 + from .logger import Experiment as Experiment + from .logger import Project as Project + from .logger import Prompt as Prompt + from .generated_types import * if os.getenv("BRAINTRUST_INSTRUMENT_THREADS", "").lower() in ("true", "1", "yes"): @@ -60,6 +71,7 @@ def is_equal(expected, output): except Exception: pass # Never break on import +from . import generated_types as _generated_types # noqa: I001 from .audit import * from .auto import auto_instrument as auto_instrument from .dataset_pipeline import * @@ -67,126 +79,14 @@ def is_equal(expected, output): from .framework2 import * from .functions.invoke import * from .functions.stream import * -from .generated_types import ( # noqa: F401 - Acl, - AclObjectType, - Agent, - AISecret, - AnyModelParams, - ApiKey, - AsyncScoringControl, - AsyncScoringState, - AttachmentReference, - AttachmentStatus, - AutomationStatus, - BatchedFacetData, - BraintrustAttachmentReference, - BraintrustModelParams, - CallEvent, - ChatCompletionContentPart, - ChatCompletionContentPartFileFile, - ChatCompletionContentPartFileWithTitle, - ChatCompletionContentPartImageWithTitle, - ChatCompletionContentPartText, - ChatCompletionContentPartTextWithTitle, - ChatCompletionMessageParam, - ChatCompletionMessageReasoning, - ChatCompletionMessageToolCall, - ChatCompletionOpenAIMessageParam, - ChatCompletionTool, - CodeBundle, - DatasetEvent, - DatasetSnapshot, - EnvVar, - ExperimentEvent, - ExtendedSavedFunctionId, - ExternalAttachmentReference, - FacetData, - FacetPreprocessorId, - Function, - FunctionData, - FunctionFormat, - FunctionId, - FunctionIdRef, - FunctionObjectType, - FunctionOutputType, - FunctionTypeEnum, - FunctionTypeEnumNullish, - GitMetadataSettings, - GraphData, - GraphEdge, - GraphNode, - Group, - GroupScope, - IfExists, - ImageRenderingMode, - InvokeFunction, - InvokeParent, - MCPServer, - MessageRole, - ModelParams, - NullableSavedFunctionId, - ObjectReference, - ObjectReferenceNullish, - OnlineScoreConfig, - Organization, - OrgAutomation, - Permission, - PreprocessorId, - ProjectAutomation, - ProjectGroup, - ProjectLogsEvent, - ProjectScore, - ProjectScoreCategories, - ProjectScoreCategory, - ProjectScoreCondition, - ProjectScoreConfig, - ProjectScoreType, - ProjectSettings, - ProjectTag, - PromptBlockData, - PromptBlockDataNullish, - PromptData, - PromptDataNullish, - PromptOptions, - PromptOptionsNullish, - PromptParserNullish, - PromptSessionEvent, - RepoInfo, - ResponseFormat, - ResponseFormatJsonSchema, - ResponseFormatNullish, - RetentionObjectType, - Role, - RunEval, - SavedFunctionId, - ServiceToken, - SpanAttributes, - SpanIFrame, - SpanScope, - SpanType, - SSEConsoleEventData, - SSEProgressEventData, - StreamingMode, - ToolFunctionDefinition, - TopicAutomationConfig, - TopicAutomationDataScope, - TopicAutomationFacetModel, - TopicDigestAutomationConfig, - TopicMapData, - TopicMapFunctionAutomation, - TopicMapGenerationSettings, - TraceScope, - TriggeredFunctionState, - UploadStatus, - User, - View, - ViewData, - ViewDataSearch, - ViewOptions, - WindowedAutomationConfig, -) -from .integrations.ai_sdk import setup_ai_sdk as setup_ai_sdk + +# Keep this before the logger wildcard so its existing runtime collision +# precedence remains unchanged while new generated names are picked up. +for _name in _generated_types.__all__: + if _name not in {"Dataset", "Experiment", "Project", "Prompt"}: + globals()[_name] = getattr(_generated_types, _name) + +from .integrations.ai_sdk import setup_ai_sdk as setup_ai_sdk # noqa: I001 from .integrations.anthropic import wrap_anthropic as wrap_anthropic from .integrations.instructor import wrap_instructor as wrap_instructor from .integrations.litellm import wrap_litellm as wrap_litellm @@ -195,6 +95,10 @@ def is_equal(expected, output): from .integrations.pydantic_ai import setup_pydantic_ai as setup_pydantic_ai from .logger import * from .logger import ( + Dataset as Dataset, + Experiment as Experiment, + Project as Project, + Prompt as Prompt, _internal_get_global_state, # noqa: F401 # type: ignore[reportUnusedImport] _internal_reset_global_state, # noqa: F401 # type: ignore[reportUnusedImport] _internal_with_custom_background_logger, # noqa: F401 # type: ignore[reportUnusedImport] diff --git a/py/src/braintrust/type_tests/test_public_exports.py b/py/src/braintrust/type_tests/test_public_exports.py index 66443861..0defbbf1 100644 --- a/py/src/braintrust/type_tests/test_public_exports.py +++ b/py/src/braintrust/type_tests/test_public_exports.py @@ -2,12 +2,17 @@ The static resource check keeps mypy from resolving generated ``TypedDict`` names instead of the public runtime classes. The runtime checks cover PEP 484 -aliasing for pyright's ``reportPrivateImportUsage`` rule. +aliasing for pyright's ``reportPrivateImportUsage`` rule and keep generated +exports synchronized with the package root. """ +import subprocess +import sys + import braintrust import pytest from braintrust import ( + Acl, auto_instrument, setup_ai_sdk, setup_pydantic_ai, @@ -34,14 +39,37 @@ def accepts_public_resource_types( dataset: braintrust.Dataset, project: braintrust.Project, prompt: braintrust.Prompt, + acl: Acl, ) -> None: experiment.fetch() dataset.fetch() _ = project.name prompt.build() + _ = acl["id"] @pytest.mark.parametrize("name,imported", _PUBLIC_SYMBOLS) def test_top_level_public_symbol(name: str, imported: object) -> None: assert callable(imported) assert callable(getattr(braintrust, name)) + + +def test_generated_exports_follow_generated_all() -> None: + script = """ +import importlib + +import braintrust +from braintrust import generated_types + +future_type = type("FutureGeneratedType", (), {}) +generated_types.FutureGeneratedType = future_type +generated_types.__all__.append("FutureGeneratedType") +try: + importlib.reload(braintrust) + assert braintrust.FutureGeneratedType is future_type +finally: + generated_types.__all__.remove("FutureGeneratedType") + del generated_types.FutureGeneratedType +""" + + subprocess.run([sys.executable, "-c", script], check=True)