use sys.stdlib_module_names in restrictions. Move dynamic import warning to only apply to non-passed through modules (#1222)

Author: VegetarianOrc

temporalio/worker/workflow_sandbox/_importer.py

@@ -201,17 +201,6 @@ def _import(
 
         # Check module restrictions and passthrough modules
         if full_name not in sys.modules:
-            # Issue a warning if appropriate
-            if (
-                self.restriction_context.in_activation
-                and self._is_import_notification_policy_applied(
-                    temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT
-                )
-            ):
-                warnings.warn(
-                    f"Module {full_name} was imported after initial workflow load."
-                )
-
             # Make sure not an entirely invalid module
             self._assert_valid_module(full_name)
 
@@ -237,6 +226,17 @@ def _import(
                     setattr(sys.modules[parent], child, sys.modules[full_name])
                 # All children of this module that are on the original sys
                 # modules but not here and are passthrough
+            else:
+                # Issue a warning if appropriate
+                if (
+                    self.restriction_context.in_activation
+                    and self._is_import_notification_policy_applied(
+                        temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_DYNAMIC_IMPORT
+                    )
+                ):
+                    warnings.warn(
+                        f"Module {full_name} was imported after initial workflow load."
+                    )
 
             # If the module is __temporal_main__ and not already in sys.modules,
             # we load it from whatever file __main__ was originally in

temporalio/worker/workflow_sandbox/_restrictions.py

@@ -34,8 +34,6 @@
     cast,
 )
 
-from typing_extensions import Self
-
 try:
     import pydantic
     import pydantic_core
@@ -529,44 +527,6 @@ def with_child_unrestricted(self, *child_path: str) -> SandboxMatcher:
     }
 )
 
-# sys.stdlib_module_names is only available on 3.10+, so we hardcode here. A
-# test will fail if this list doesn't match the latest Python version it was
-# generated against, spitting out the expected list. This is a string instead
-# of a list of strings due to black wanting to format this to one item each
-# line in a list.
-_stdlib_module_names = (
-    "__future__,_abc,_aix_support,_ast,_asyncio,_bisect,_blake2,_bootsubprocess,_bz2,_codecs,"
-    "_codecs_cn,_codecs_hk,_codecs_iso2022,_codecs_jp,_codecs_kr,_codecs_tw,_collections,"
-    "_collections_abc,_compat_pickle,_compression,_contextvars,_crypt,_csv,_ctypes,_curses,"
-    "_curses_panel,_datetime,_dbm,_decimal,_elementtree,_frozen_importlib,_frozen_importlib_external,"
-    "_functools,_gdbm,_hashlib,_heapq,_imp,_io,_json,_locale,_lsprof,_lzma,_markupbase,"
-    "_md5,_msi,_multibytecodec,_multiprocessing,_opcode,_operator,_osx_support,_overlapped,"
-    "_pickle,_posixshmem,_posixsubprocess,_py_abc,_pydecimal,_pyio,_queue,_random,_scproxy,"
-    "_sha1,_sha256,_sha3,_sha512,_signal,_sitebuiltins,_socket,_sqlite3,_sre,_ssl,_stat,"
-    "_statistics,_string,_strptime,_struct,_symtable,_thread,_threading_local,_tkinter,"
-    "_tokenize,_tracemalloc,_typing,_uuid,_warnings,_weakref,_weakrefset,_winapi,_zoneinfo,"
-    "abc,aifc,antigravity,argparse,array,ast,asynchat,asyncio,asyncore,atexit,audioop,"
-    "base64,bdb,binascii,bisect,builtins,bz2,cProfile,calendar,cgi,cgitb,chunk,cmath,cmd,"
-    "code,codecs,codeop,collections,colorsys,compileall,concurrent,configparser,contextlib,"
-    "contextvars,copy,copyreg,crypt,csv,ctypes,curses,dataclasses,datetime,dbm,decimal,"
-    "difflib,dis,distutils,doctest,email,encodings,ensurepip,enum,errno,faulthandler,fcntl,"
-    "filecmp,fileinput,fnmatch,fractions,ftplib,functools,gc,genericpath,getopt,getpass,"
-    "gettext,glob,graphlib,grp,gzip,hashlib,heapq,hmac,html,http,idlelib,imaplib,imghdr,"
-    "imp,importlib,inspect,io,ipaddress,itertools,json,keyword,lib2to3,linecache,locale,"
-    "logging,lzma,mailbox,mailcap,marshal,math,mimetypes,mmap,modulefinder,msilib,msvcrt,"
-    "multiprocessing,netrc,nis,nntplib,nt,ntpath,nturl2path,numbers,opcode,operator,optparse,"
-    "os,ossaudiodev,pathlib,pdb,pickle,pickletools,pipes,pkgutil,platform,plistlib,poplib,"
-    "posix,posixpath,pprint,profile,pstats,pty,pwd,py_compile,pyclbr,pydoc,pydoc_data,"
-    "pyexpat,queue,quopri,random,re,readline,reprlib,resource,rlcompleter,runpy,sched,"
-    "secrets,select,selectors,shelve,shlex,shutil,signal,site,smtpd,smtplib,sndhdr,socket,"
-    "socketserver,spwd,sqlite3,sre_compile,sre_constants,sre_parse,ssl,stat,statistics,"
-    "string,stringprep,struct,subprocess,sunau,symtable,sys,sysconfig,syslog,tabnanny,"
-    "tarfile,telnetlib,tempfile,termios,textwrap,this,threading,time,timeit,tkinter,token,"
-    "tokenize,tomllib,trace,traceback,tracemalloc,tty,turtle,turtledemo,types,typing,unicodedata,"
-    "unittest,urllib,uu,uuid,venv,warnings,wave,weakref,webbrowser,winreg,winsound,wsgiref,"
-    "xdrlib,xml,xmlrpc,zipapp,zipfile,zipimport,zlib,zoneinfo"
-)
-
 SandboxRestrictions.passthrough_modules_maximum = (
     SandboxRestrictions.passthrough_modules_with_temporal
     | {
@@ -575,7 +535,7 @@ def with_child_unrestricted(self, *child_path: str) -> SandboxMatcher:
         # manually setting sys.modules["os.path"]) they have certain child
         # expectations.
         v
-        for v in _stdlib_module_names.split(",")
+        for v in sys.stdlib_module_names
         if v != "sys"
     }
 )

tests/worker/workflow_sandbox/test_restrictions.py

@@ -13,7 +13,6 @@
     SandboxMatcher,
     SandboxRestrictions,
     _RestrictedProxy,
-    _stdlib_module_names,
 )
 
 
@@ -31,8 +30,8 @@ def test_workflow_sandbox_stdlib_module_names():
         code_lines[-1] += mod_name
     code = '_stdlib_module_names = (\n    "' + '"\n    "'.join(code_lines) + '"\n)'
     # TODO(cretz): Point releases may add modules :-(
-    assert (
-        actual_names == _stdlib_module_names
+    assert actual_names == ",".join(
+        sys.stdlib_module_names
     ), f"Expecting names as {actual_names}. In code as:\n{code}"
 
 

tests/worker/workflow_sandbox/test_runner.py

@@ -602,27 +602,19 @@ async def test_workflow_sandbox_import_errors(client: Client):
         workflows=[LazyImportWorkflow],
         workflow_runner=SandboxedWorkflowRunner(restrictions),
     ) as worker:
-        with pytest.warns() as recorder:
-            with pytest.raises(WorkflowFailureError) as err:
-                await client.execute_workflow(
-                    LazyImportWorkflow.run,
-                    id=f"workflow-{uuid.uuid4()}",
-                    task_queue=worker.task_queue,
-                )
-
-            assert isinstance(err.value.cause, ApplicationError)
-            assert err.value.cause.type == "UnintentionalPassthroughError"
-            assert (
-                "Module tests.worker.workflow_sandbox.testmodules.lazy_module was not intentionally passed through to the sandbox."
-                == err.value.cause.message
+        with pytest.raises(WorkflowFailureError) as err:
+            await client.execute_workflow(
+                LazyImportWorkflow.run,
+                id=f"workflow-{uuid.uuid4()}",
+                task_queue=worker.task_queue,
             )
 
-            _assert_expected_warnings(
-                recorder,
-                {
-                    "Module tests.worker.workflow_sandbox.testmodules.lazy_module was imported after initial workflow load.",
-                },
-            )
+        assert isinstance(err.value.cause, ApplicationError)
+        assert err.value.cause.type == "UnintentionalPassthroughError"
+        assert (
+            "Module tests.worker.workflow_sandbox.testmodules.lazy_module was not intentionally passed through to the sandbox."
+            == err.value.cause.message
+        )
 
 
 @workflow.defn