Merge pull request #1328 from moreati/issue1325-scan_code_imports-refactor

mitogen: Refactor `mitogen.master.scan_code_imports()` -> `mitogen.imports.codeobj_imports()`

Author: moreati

ansible_mitogen/module_finder.py

@@ -44,7 +44,7 @@
     # Python < 3.4, PEP 302 Import Hooks
     import imp
 
-import mitogen.master
+import mitogen.imports
 
 
 LOG = logging.getLogger(__name__)
@@ -146,7 +146,7 @@ def scan_fromlist(code):
     >>> list(scan_fromlist(code))
     [(0, 'a'), (0, 'b.c'), (0, 'd.e.f'), (0, 'g.h'), (0, 'g.i')]
     """
-    for level, modname_s, fromlist in mitogen.master.scan_code_imports(code):
+    for level, modname_s, fromlist in mitogen.imports.codeobj_imports(code):
         for name in fromlist:
             yield level, str('%s.%s' % (modname_s, name))
         if not fromlist:
@@ -172,7 +172,7 @@ def walk_imports(code, prefix=None):
         prefix = ''
     pattern = re.compile(r'(^|\.)(\w+)')
     start = len(prefix)
-    for _, name, fromlist in mitogen.master.scan_code_imports(code):
+    for _, name, fromlist in mitogen.imports.codeobj_imports(code):
         if not name.startswith(prefix):
             continue
         for match in pattern.finditer(name, start):

docs/changelog.rst

@@ -21,8 +21,12 @@ To avail of fixes in an unreleased version, please download a ZIP file
 In progress (unreleased)
 ------------------------
 
+* :gh:issue:`1325` :mod:`mitogen`: Refactor
+  ``mitogen.master.scan_code_imports()`` as
+  :func:`mitogen.import.codeobj_imports` and speed-up by 1.5 - 2.5 x
 * :gh:issue:`1329` CI: Refactor and de-duplicate Github Actions workflow
 * :gh:issue:`1315` CI: macOS: Increase failed logins limit of test users
+* :gh:issue:`1325` tests: Improve ``master_test.ScanCodeImportsTest`` coverage
 
 
 v0.3.26 (2025-08-04)

mitogen/core.py

@@ -1300,6 +1300,7 @@ class Importer(object):
         'kubectl',
         'fakessh',
         'fork',
+        'imports',
         'jail',
         'lxc',
         'lxd',

mitogen/imports/__init__.py

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: 2025 Mitogen authors <https://github.com/mitogen-hq>
+# SPDX-License-Identifier: MIT
+# !mitogen: minify_safe
+
+import sys
+
+if sys.version_info >= (3, 6):
+    from mitogen.imports._py36 import _code_imports
+elif sys.version_info >= (2, 5):
+    from mitogen.imports._py2 import _code_imports_py25 as _code_imports
+else:
+    from mitogen.imports._py2 import _code_imports_py24 as _code_imports
+
+
+def codeobj_imports(co):
+    """
+    Yield (level, modname, names) tuples by scanning the code object `co`.
+
+    Top level `import mod` & `from mod import foo` statements are matched.
+    Those inside a `class ...` or `def ...` block are currently skipped.
+
+    >>> co = compile('import a, b; from c import d, e as f', '<str>', 'exec')
+    >>> list(codeobj_imports(co))  # doctest: +ELLIPSIS
+    [(..., 'a', ()), (..., 'b', ()), (..., 'c', ('d', 'e'))]
+
+    :return:
+        Generator producing `(level, modname, names)` tuples, where:
+
+        * `level`:
+            -1 implicit relative (Python 2.x default)
+            0  absolute (Python 3.x, `from __future__ import absolute_import`)
+            >0 explicit relative (`from . import a`, `from ..b, import c`)
+        * `modname`: Name of module to import, or to import `names` from.
+        * `names`: tuple of names in `from mod import ..`.
+    """
+    return _code_imports(co.co_code, co.co_consts, co.co_names)

mitogen/imports/_py2.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: 2025 Mitogen authors <https://github.com/mitogen-hq>
+# SPDX-License-Identifier: MIT
+# !mitogen: minify_safe
+
+import array
+import itertools
+import opcode
+
+
+IMPORT_NAME = opcode.opmap['IMPORT_NAME']
+LOAD_CONST = opcode.opmap['LOAD_CONST']
+
+
+def _opargs(code, _have_arg=opcode.HAVE_ARGUMENT):
+    it = iter(array.array('B', code))
+    nexti = it.next
+    for i in it:
+        if i >= _have_arg:
+            yield (i, nexti() | (nexti() << 8))
+        else:
+            yield (i, None)
+
+
+def _code_imports_py25(code, consts, names):
+    it1, it2, it3 = itertools.tee(_opargs(code), 3)
+    try:
+        next(it2)
+        next(it3)
+        next(it3)
+    except StopIteration:
+        return
+    for oparg1, oparg2, (op3, arg3) in itertools.izip(it1, it2, it3):
+        if op3 != IMPORT_NAME:
+            continue
+        op1, arg1 = oparg1
+        op2, arg2 = oparg2
+        if op1 != LOAD_CONST or op2 != LOAD_CONST:
+            continue
+        yield (consts[arg1], names[arg3], consts[arg2] or ())
+
+
+def _code_imports_py24(code, consts, names):
+    it1, it2 = itertools.tee(_opargs(code), 2)
+    try:
+        next(it2)
+    except StopIteration:
+        return
+    for oparg1, (op2, arg2) in itertools.izip(it1, it2):
+        if op2 != IMPORT_NAME:
+            continue
+        op1, arg1 = oparg1
+        if op1 != LOAD_CONST:
+            continue
+        yield (-1, names[arg2], consts[arg1] or ())

mitogen/imports/_py36.py

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: 2025 Mitogen authors <https://github.com/mitogen-hq>
+# SPDX-License-Identifier: MIT
+# !mitogen: minify_safe
+
+import opcode
+
+IMPORT_NAME = opcode.opmap['IMPORT_NAME']
+LOAD_CONST = opcode.opmap['LOAD_CONST']
+
+
+def _code_imports(code, consts, names):
+    start = 4
+    while True:
+        op3_idx = code.find(IMPORT_NAME, start, -1)
+        if op3_idx < 0:
+            return
+        if op3_idx % 2:
+            start = op3_idx + 1
+            continue
+        if code[op3_idx-4] != LOAD_CONST or code[op3_idx-2] != LOAD_CONST:
+            start = op3_idx + 2
+            continue
+        start = op3_idx + 6
+        arg1, arg2, arg3 = code[op3_idx-3], code[op3_idx-1], code[op3_idx+1]
+        yield (consts[arg1], names[arg3], consts[arg2] or ())

mitogen/master.py

@@ -35,10 +35,8 @@
 contexts.
 """
 
-import dis
 import errno
 import inspect
-import itertools
 import logging
 import os
 import pkgutil
@@ -83,21 +81,18 @@ def _find_loader(fullname):
 
 import mitogen
 import mitogen.core
+import mitogen.imports
 import mitogen.minify
 import mitogen.parent
 
 from mitogen.core import any
 from mitogen.core import b
 from mitogen.core import IOLOG
 from mitogen.core import LOG
-from mitogen.core import next
 from mitogen.core import str_partition
 from mitogen.core import str_rpartition
 from mitogen.core import to_text
 
-imap = getattr(itertools, 'imap', map)
-izip = getattr(itertools, 'izip', zip)
-
 RLOG = logging.getLogger('mitogen.ctx')
 
 
@@ -253,80 +248,6 @@ def _get_core_source():
     mitogen.parent._get_core_source = _get_core_source
 
 
-LOAD_CONST = dis.opname.index('LOAD_CONST')
-IMPORT_NAME = dis.opname.index('IMPORT_NAME')
-
-
-def _getarg(nextb, c):
-    if c >= dis.HAVE_ARGUMENT:
-        return nextb() | (nextb() << 8)
-
-
-if sys.version_info < (3, 0):
-    def iter_opcodes(co):
-        # Yield `(op, oparg)` tuples from the code object `co`.
-        ordit = imap(ord, co.co_code)
-        nextb = ordit.next
-        return ((c, _getarg(nextb, c)) for c in ordit)
-elif sys.version_info < (3, 6):
-    def iter_opcodes(co):
-        # Yield `(op, oparg)` tuples from the code object `co`.
-        ordit = iter(co.co_code)
-        nextb = ordit.__next__
-        return ((c, _getarg(nextb, c)) for c in ordit)
-else:
-    def iter_opcodes(co):
-        # Yield `(op, oparg)` tuples from the code object `co`.
-        ordit = iter(co.co_code)
-        nextb = ordit.__next__
-        # https://github.com/abarnert/cpython/blob/c095a32f/Python/wordcode.md
-        return ((c, nextb()) for c in ordit)
-
-
-def scan_code_imports(co):
-    """
-    Given a code object `co`, scan its bytecode yielding any ``IMPORT_NAME``
-    and associated prior ``LOAD_CONST`` instructions representing an `Import`
-    statement or `ImportFrom` statement.
-
-    :return:
-        Generator producing `(level, modname, namelist)` tuples, where:
-
-        * `level`: -1 for normal import, 0, for absolute import, and >0 for
-          relative import.
-        * `modname`: Name of module to import, or from where `namelist` names
-          are imported.
-        * `namelist`: for `ImportFrom`, the list of names to be imported from
-          `modname`.
-    """
-    opit = iter_opcodes(co)
-    opit, opit2, opit3 = itertools.tee(opit, 3)
-
-    try:
-        next(opit2)
-        next(opit3)
-        next(opit3)
-    except StopIteration:
-        return
-
-    if sys.version_info >= (2, 5):
-        for oparg1, oparg2, (op3, arg3) in izip(opit, opit2, opit3):
-            if op3 == IMPORT_NAME:
-                op2, arg2 = oparg2
-                op1, arg1 = oparg1
-                if op1 == op2 == LOAD_CONST:
-                    yield (co.co_consts[arg1],
-                           co.co_names[arg3],
-                           co.co_consts[arg2] or ())
-    else:
-        # Python 2.4 did not yet have 'level', so stack format differs.
-        for oparg1, (op2, arg2) in izip(opit, opit2):
-            if op2 == IMPORT_NAME:
-                op1, arg1 = oparg1
-                if op1 == LOAD_CONST:
-                    yield (-1, co.co_names[arg2], co.co_consts[arg1] or ())
-
-
 class ThreadWatcher(object):
     """
     Manage threads that wait for another thread to shut down, before invoking
@@ -1029,7 +950,7 @@ def find_related_imports(self, fullname):
         maybe_names = list(self.generate_parent_names(fullname))
 
         co = compile(src, modpath, 'exec')
-        for level, modname, namelist in scan_code_imports(co):
+        for level, modname, namelist in mitogen.imports.codeobj_imports(co):
             if level == -1:
                 modnames = [modname, '%s.%s' % (fullname, modname)]
             else: