Do some renaming, apply linter.

Author: VegetarianOrc

temporalio/worker/workflow_sandbox/__init__.py

@@ -55,10 +55,10 @@
 #   * https://github.com/GrahamDumpleton/wrapt/issues/130
 
 from ._restrictions import (
+    DisallowedUnintentionalPassthroughError,
     RestrictedWorkflowAccessError,
     SandboxMatcher,
     SandboxRestrictions,
-    DisallowedUnintentionalPassthroughError,
 )
 from ._runner import SandboxedWorkflowRunner
 

temporalio/worker/workflow_sandbox/_importer.py

@@ -38,9 +38,9 @@
 import temporalio.workflow
 
 from ._restrictions import (
+    DisallowedUnintentionalPassthroughError,
     RestrictedModule,
     RestrictedWorkflowAccessError,
-    DisallowedUnintentionalPassthroughError,
     RestrictionContext,
     SandboxRestrictions,
 )
@@ -205,7 +205,7 @@ def _import(
             if (
                 self.restriction_context.in_activation
                 and self._is_import_policy_applied(
-                    temporalio.workflow.SandboxImportPolicy.WARN_ON_DYNAMIC_IMPORT
+                    temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT
                 )
             ):
                 warnings.warn(
@@ -295,13 +295,15 @@ def module_configured_passthrough(self, name: str) -> bool:
         return True
 
     def _is_import_policy_applied(
-        self, policy: temporalio.workflow.SandboxImportPolicy
+        self, policy: temporalio.workflow.SandboxImportNotificationPolicy
     ) -> bool:
-        applied_policy = temporalio.workflow.unsafe.current_sandbox_import_policy()
-        if applied_policy != temporalio.workflow.SandboxImportPolicy.UNSET:
+        applied_policy = (
+            temporalio.workflow.unsafe.current_sandbox_import_notification_policy()
+        )
+        if applied_policy:
             return policy in applied_policy
 
-        return policy in self.restrictions.import_policy
+        return policy in self.restrictions.import_notification_policy
 
     def _maybe_passthrough_module(self, name: str) -> Optional[types.ModuleType]:
         # If imports not passed through and all modules are not passed through
@@ -311,13 +313,12 @@ def _maybe_passthrough_module(self, name: str) -> Optional[types.ModuleType]:
             and not self.module_configured_passthrough(name)
         ):
             if self._is_import_policy_applied(
-                temporalio.workflow.SandboxImportPolicy.RAISE_ON_NON_PASSTHROUGH
+                temporalio.workflow.SandboxImportNotificationPolicy.RAISE_ON_UNINTENTIONAL_PASSTHROUGH
             ):
-                # TODO(amazzeo): this is not an appropriate error type
                 raise DisallowedUnintentionalPassthroughError(name)
 
             if self._is_import_policy_applied(
-                temporalio.workflow.SandboxImportPolicy.WARN_ON_NON_PASSTHROUGH
+                temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH
             ):
                 warnings.warn(
                     f"Module {name} was not intentionally passed through to the sandbox."

temporalio/worker/workflow_sandbox/_restrictions.py

@@ -84,9 +84,17 @@ def default_message(qualified_name: str) -> str:
 
 # TODO(amazzeo): is NondeterminisimError appropriate as a subclass?
 class DisallowedUnintentionalPassthroughError(temporalio.workflow.NondeterminismError):
+    """Error that occurs when a workflow unintentionally passes an import to the sandbox when
+    the import notification policy includes :py:attr:`temporalio.workflow.SandboxImportNotificationPolicy.RAISE_ON_NON_PASSTHROUGH`.
+
+    Attributes:
+        qualified_name: Fully qualified name of what was passed through to the sandbox.
+    """
+
     def __init__(
         self, qualified_name: str, *, override_message: Optional[str] = None
     ) -> None:
+        """Create a disallowed unintentional passthrough error."""
         super().__init__(
             override_message
             or DisallowedUnintentionalPassthroughError.default_message(qualified_name)
@@ -95,7 +103,7 @@ def __init__(
 
     @staticmethod
     def default_message(qualified_name: str) -> str:
-        """Get default message for restricted access."""
+        """Get default message for disallowed unintential passthrough."""
         return f"Module {qualified_name} was not intentionally passed through to the sandbox."
 
 
@@ -127,14 +135,43 @@ class methods (including __init__, etc). The check compares the against the
     fully qualified path to the item.
     """
 
-    import_policy: temporalio.workflow.SandboxImportPolicy = (
-        temporalio.workflow.SandboxImportPolicy.WARN_ON_DYNAMIC_IMPORT
-    )
+    import_notification_policy: temporalio.workflow.SandboxImportNotificationPolicy
+    """
+    The import notification policy to use when an import is triggered during workflow loading or execution. See :py:class:`temporalio.workflow.SandboxImportNotificationPolicy` for options.
+    """
+
+    import_notification_policy_default: ClassVar[
+        temporalio.workflow.SandboxImportNotificationPolicy
+    ]
+    """
+    The default import notification policy. Equivalent to :py:attr:`temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT`.
+    """
+
+    import_notification_policy_all_warnings: ClassVar[
+        temporalio.workflow.SandboxImportNotificationPolicy
+    ]
+    """
+    An import notification policy that enables warnings on dynamic imports during a 
+    workflow and imports that are unintentionally passed through to the sandbox.
+    Equivalent to :py:attr:`temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT` | :py:attr:`temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH` 
+    """
 
-    import_policy_all_warnings: ClassVar[temporalio.workflow.SandboxImportPolicy]
-    import_policy_disallow_unintentional_passthrough: ClassVar[
-        temporalio.workflow.SandboxImportPolicy
+    import_notification_policy_disallow_unintentional_passthrough: ClassVar[
+        temporalio.workflow.SandboxImportNotificationPolicy
     ]
+    """
+    An import notification policy that enables warnings on dynamic imports during a 
+    workflow and raises an error when imports are unintentionally passed through to the sandbox.
+    Equivalent to :py:attr:`temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT` | :py:attr:`temporalio.workflow.SandboxImportNotificationPolicy.RAISE_ON_UNINTENTIONAL_PASSTHROUGH` 
+    """
+
+    import_notification_policy_silent: ClassVar[
+        temporalio.workflow.SandboxImportNotificationPolicy
+    ]
+    """
+    An import notification policy that disables all warnings and errors.
+    Equivalent to :py:attr:`temporalio.workflow.SandboxImportNotificationPolicy.SILENT`.
+    """
 
     passthrough_all_modules: bool = False
     """
@@ -196,6 +233,12 @@ def with_passthrough_all_modules(self) -> SandboxRestrictions:
         """
         return dataclasses.replace(self, passthrough_all_modules=True)
 
+    def with_import_policy(
+        self, policy: temporalio.workflow.SandboxImportNotificationPolicy
+    ) -> SandboxRestrictions:
+        """Create a new restriction set with the given import notification policy as the :py:attr:`import_policy`."""
+        return dataclasses.replace(self, import_notification_policy=policy)
+
 
 # We intentionally use specific fields instead of generic "matcher" callbacks
 # for optimization reasons.
@@ -311,7 +354,6 @@ def access_matcher(
         Returns:
             The matcher if matched.
         """
-
         # We prefer to avoid recursion
         matcher = self
         for v in child_path:
@@ -525,13 +567,19 @@ def with_child_unrestricted(self, *child_path: str) -> SandboxMatcher:
     }
 )
 
-SandboxRestrictions.import_policy_all_warnings = (
-    temporalio.workflow.SandboxImportPolicy.WARN_ON_DYNAMIC_IMPORT
-    | temporalio.workflow.SandboxImportPolicy.WARN_ON_NON_PASSTHROUGH
+SandboxRestrictions.import_notification_policy_default = (
+    temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT
+)
+SandboxRestrictions.import_notification_policy_all_warnings = (
+    temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT
+    | temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH
+)
+SandboxRestrictions.import_notification_policy_disallow_unintentional_passthrough = (
+    temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT
+    | temporalio.workflow.SandboxImportNotificationPolicy.RAISE_ON_UNINTENTIONAL_PASSTHROUGH
 )
-SandboxRestrictions.import_policy_disallow_unintentional_passthrough = (
-    temporalio.workflow.SandboxImportPolicy.WARN_ON_DYNAMIC_IMPORT
-    | temporalio.workflow.SandboxImportPolicy.RAISE_ON_NON_PASSTHROUGH
+SandboxRestrictions.import_notification_policy_silent = (
+    temporalio.workflow.SandboxImportNotificationPolicy.SILENT
 )
 
 # sys.stdlib_module_names is only available on 3.10+, so we hardcode here. A
@@ -836,6 +884,7 @@ def _public_callables(parent: Any, *, exclude: Set[str] = set()) -> Set[str]:
     passthrough_modules=SandboxRestrictions.passthrough_modules_default,
     invalid_modules=SandboxMatcher.none,
     invalid_module_members=SandboxRestrictions.invalid_module_members_default,
+    import_notification_policy=SandboxRestrictions.import_notification_policy_default,
 )
 
 

temporalio/workflow.py

@@ -1410,24 +1410,20 @@ async def wait_condition(
 _sandbox_unrestricted = threading.local()
 _in_sandbox = threading.local()
 _imports_passed_through = threading.local()
+_sandbox_import_notification_policy_override = threading.local()
 
 
-# TODO(amazzeo): this name is confusing with existing import passthrough phrasing.
-class SandboxImportPolicy(Flag):
-    # TODO(amazzeo): Need to establish consistent phrasing about the time this is applied
-    """Defines the behavior taken when modules are imported into the sandbox after the workflow is initially loaded."""
+class SandboxImportNotificationPolicy(Flag):
+    """Defines the behavior taken when modules are imported into the sandbox after the workflow is initially loaded or unintentionally missing from the passthrough list."""
 
-    UNSET = auto()
-    """All imports that do not violate sandbox restrictions are allowed and no warning is generated."""
+    SILENT = auto()
+    """Allow imports that do not violate sandbox restrictions and no warnings are generated."""
     WARN_ON_DYNAMIC_IMPORT = auto()
-    """Issue a warning when an import is triggered in the sandbox after initial workflow load."""
-    WARN_ON_NON_PASSTHROUGH = auto()
-    """Issue a warning when an import is triggered in the sandbox that was not passed through."""
-    RAISE_ON_NON_PASSTHROUGH = auto()
-    """Raise an error when an import is triggered in the sandbox that was not passed through."""
-
-
-_sandbox_import_policy = threading.local()
+    """Allows dynamic imports that do not violate sandbox restrictions but issues a warning when an import is triggered in the sandbox after initial workflow load."""
+    WARN_ON_UNINTENTIONAL_PASSTHROUGH = auto()
+    """Allows imports that do not violate sandbox restrictions but issues a warning when an import is triggered in the sandbox that was unintentionally passed through."""
+    RAISE_ON_UNINTENTIONAL_PASSTHROUGH = auto()
+    """Raise an error when an import is triggered in the sandbox that was unintentionally passed through."""
 
 
 class unsafe:
@@ -1517,24 +1513,33 @@ def imports_passed_through() -> Iterator[None]:
             _imports_passed_through.value = False
 
     @staticmethod
-    def current_sandbox_import_policy() -> SandboxImportPolicy:
+    def current_sandbox_import_notification_policy() -> (
+        Optional[SandboxImportNotificationPolicy]
+    ):
+        """Gets the current import notification policy override if one is set."""
         applied_policy = getattr(
-            _sandbox_import_policy, "value", SandboxImportPolicy.UNSET
+            _sandbox_import_notification_policy_override,
+            "value",
+            None,
         )
         return applied_policy
 
     @staticmethod
     @contextmanager
-    def sandbox_import_policy(policy: SandboxImportPolicy) -> Iterator[None]:
-        # TODO(amazzeo): the default behavior here seems inappropriate
-        original_policy = _sandbox_import_policy.value = getattr(
-            _sandbox_import_policy, "value", SandboxImportPolicy.UNSET
+    def sandbox_import_notification_policy(
+        policy: SandboxImportNotificationPolicy,
+    ) -> Iterator[None]:
+        """Context manager to apply the given import notification policy."""
+        original_policy = _sandbox_import_notification_policy_override.value = getattr(
+            _sandbox_import_notification_policy_override,
+            "value",
+            None,
         )
-        _sandbox_import_policy.value = policy
+        _sandbox_import_notification_policy_override.value = policy
         try:
             yield None
         finally:
-            _sandbox_import_policy.value = original_policy
+            _sandbox_import_notification_policy_override.value = original_policy
 
 
 class LoggerAdapter(logging.LoggerAdapter):

tests/worker/workflow_sandbox/test_runner.py

@@ -25,11 +25,11 @@
     WorkflowInterceptorClassInput,
 )
 from temporalio.worker.workflow_sandbox import (
+    DisallowedUnintentionalPassthroughError,
     RestrictedWorkflowAccessError,
     SandboxedWorkflowRunner,
     SandboxMatcher,
     SandboxRestrictions,
-    DisallowedUnintentionalPassthroughError,
 )
 from tests.helpers import assert_eq_eventually
 from tests.worker.workflow_sandbox.testmodules import stateful_module
@@ -41,7 +41,7 @@
 _ = os.name
 
 # This used to fail because our __init__ couldn't handle metaclass init
-import zipfile
+import zipfile  # noqa: E402
 
 
 class MyZipFile(zipfile.ZipFile):
@@ -495,10 +495,10 @@ class _TestWorkflowInboundInterceptor(WorkflowInboundInterceptor):
     async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any:
         # import in the interceptor to show it will be captured
         # applying this policy should squelch the "after initial workload" warning
-        with workflow.unsafe.sandbox_import_policy(
-            workflow.SandboxImportPolicy.WARN_ON_NON_PASSTHROUGH
+        with workflow.unsafe.sandbox_import_notification_policy(
+            workflow.SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH
         ):
-            import tests.worker.workflow_sandbox.testmodules.lazy_module_interceptor
+            import tests.worker.workflow_sandbox.testmodules.lazy_module_interceptor  # noqa: F401
 
         return await super().execute_workflow(input)
 
@@ -515,7 +515,7 @@ class LazyImportWorkflow:
     @workflow.run
     async def run(self) -> None:
         try:
-            import tests.worker.workflow_sandbox.testmodules.lazy_module
+            import tests.worker.workflow_sandbox.testmodules.lazy_module  # noqa: F401
         except DisallowedUnintentionalPassthroughError as err:
             raise ApplicationError(
                 str(err), type="DisallowedUnintentionalPassthroughError"
@@ -554,7 +554,7 @@ async def test_workflow_sandbox_import_default_warnings(client: Client):
 async def test_workflow_sandbox_import_all_warnings(client: Client):
     restrictions = dataclasses.replace(
         SandboxRestrictions.default,
-        import_policy=SandboxRestrictions.import_policy_all_warnings,
+        import_notification_policy=SandboxRestrictions.import_notification_policy_all_warnings,
         # passthrough this test module to avoid a ton of noisy warnings
         passthrough_modules=SandboxRestrictions.passthrough_modules_default
         | {"tests.worker.workflow_sandbox.test_runner"},
@@ -587,7 +587,7 @@ async def test_workflow_sandbox_import_all_warnings(client: Client):
 async def test_workflow_sandbox_import_errors(client: Client):
     restrictions = dataclasses.replace(
         SandboxRestrictions.default,
-        import_policy=SandboxRestrictions.import_policy_disallow_unintentional_passthrough,
+        import_notification_policy=SandboxRestrictions.import_notification_policy_disallow_unintentional_passthrough,
         # passthrough this test module to avoid a ton of noisy warnings
         passthrough_modules=SandboxRestrictions.passthrough_modules_default
         | {"tests.worker.workflow_sandbox.test_runner"},
@@ -622,6 +622,42 @@ async def test_workflow_sandbox_import_errors(client: Client):
             )
 
 
+@workflow.defn
+class SupressWarningsLazyImportWorkflow:
+    @workflow.run
+    async def run(self) -> None:
+        with workflow.unsafe.sandbox_import_notification_policy(
+            SandboxRestrictions.import_notification_policy_silent
+        ):
+            try:
+                import tests.worker.workflow_sandbox.testmodules.lazy_module  # noqa: F401
+            except UserWarning:
+                raise ApplicationError("No warnings were expected")
+
+
+async def test_workflow_sandbox_import_suppress_warnings(client: Client):
+    restrictions = dataclasses.replace(
+        SandboxRestrictions.default,
+        # passthrough this test module to avoid a ton of noisy warnings
+        passthrough_modules=SandboxRestrictions.passthrough_modules_default
+        | {"tests.worker.workflow_sandbox.test_runner"},
+    )
+
+    async with Worker(
+        client,
+        task_queue=str(uuid.uuid4()),
+        workflows=[SupressWarningsLazyImportWorkflow],
+        workflow_runner=SandboxedWorkflowRunner(restrictions),
+    ) as worker:
+        with pytest.warns(None) as recorder:  # type:ignore
+            await client.execute_workflow(
+                SupressWarningsLazyImportWorkflow.run,
+                id=f"workflow-{uuid.uuid4()}",
+                task_queue=worker.task_queue,
+            )
+            assert len(recorder) == 0, "Expected no warnings to be issued"
+
+
 def _assert_expected_warnings(
     recorder: pytest.WarningsRecorder, expected_warnings: Set[str]
 ):

tests/worker/workflow_sandbox/testmodules/__init__.py

@@ -14,4 +14,5 @@
         "tests.worker.workflow_sandbox.testmodules.invalid_module_members".split("."),
         SandboxMatcher(use={"invalid_function"}),
     ),
+    import_notification_policy=SandboxRestrictions.import_notification_policy_default,
 )