Identifying Python module dependencies

This post is also available as an interactive notebook.

Background

Consider the following problem: you’d like to enable users to automatically extract a function from a Jupyter notebook and publish it as a service. Actually serializing a closure from a function in a given environment is not difficult, given the cloudpickle module. But merely constructing a serialized closure isn’t enough, since in general this function may require other modules to be available to run in another context.

Therefore, we need some way to identify the modules required by a function (and, ultimately, the packages that provide these modules). Since engineering time is limited, it’s probably better to have an optimistic-but-incomplete (or unsound) estimate and allow users to override it (by supplying additional dependencies when they publish a function) than it is to have a sound and conservative module list.¹

The modulefinder module in the Python standard library might initially seem like an attractive option, but it is unsuitable because it operates at the level of scripts. In order to use modulefinder on a single function from a notebook, we’d either have an imprecise module list (due to running the whole notebook) or we’d need to essentially duplicate a lot of its effort in order to slice backwards from the function invocation so we could extract a suitably pruned script.

Fortunately, you can interrogate nearly any property of any object in a Python program, including functions. If we could inspect the captured variables in a closure, we could identify the ones that are functions and figure out which modules they were declared in. That would look something like this:

import inspect

def module_frontier(f):
  worklist = [f]
  seen = set()
  mods = set()
  for fn in worklist:
    cvs = inspect.getclosurevars(fn)
    gvars = cvs.globals
    for k, v in gvars.items():
      if inspect.ismodule(v):
        mods.add(v.__name__)
      elif inspect.isfunction(v) and id(v) not in seen:
        seen.add(id(v))
        mods.add(v.__module__)
        worklist.append(v)
      elif hasattr(v, "__module__"):
        mods.add(v.__module__)
  return list(mods)

The inspect module provides a friendly interface to inspecting object metadata. In the above function, we’re constructing a worklist of all of the captured variables in a given function’s closure. We’re then constructing a set of all of the modules directly or transitively referred to by those captured variables, whether these are modules referred to directly, modules declaring functions referred to by captured variables, or modules declaring other values referred to by captured variables (e.g., native functions). Note that we add any functions we find to the worklist (although we don’t handle eval or other techniques), so we’ll capture at least some of the transitive call graph in this case.

This approach seems to work pretty sensibly on simple examples:

import numpy as np
from numpy import dot

def f(a, b):
    return np.dot(a, b)

def g(a, b):
    return dot(a, b)

def h(a, b):
    return f(a, b)

{k.__name__ : module_frontier(k) for k in [f,g,h]}

{'f': ['numpy.core.multiarray', 'numpy'],
 'g': ['numpy.core.multiarray'],
 'h': ['numpy.core.multiarray', 'numpy', '__main__']}

It also works on itself, which is a relief:

module_frontier(module_frontier)

['inspect']

Problem cases

While these initial experiments are promising, we shouldn’t expect that a simple approach will cover everything we might want to do. Let’s look at a (slightly) more involved example to see if it breaks down.

We’ll use the k-means clustering implementation from scikit-learn to optimize some cluster centers in a model object. We’ll then capture that model object in a closure and analyze it to see what we might need to import to run it elsewhere.

from sklearn.cluster import KMeans
import numpy as np

data = np.random.rand(1000, 2)

model = KMeans(random_state=0).fit(data)

def km_predict_one(sample):
    sample = np.array(sample).reshape(1,-1)
    return model.predict(sample)[0]

km_predict_one([0.5, 0.5])

7

module_frontier(km_predict_one)

['sklearn.cluster.k_means_', 'numpy']

List comprehensions

So far, so good. Let’s say we want to publish this simple model as a lighter-weight service (without a scikit-learn dependency). We can get that by reimplementing the predict method from the k-means model:

centers = model.cluster_centers_
from numpy.linalg import norm

def km_predict_two(sample):
    _, idx = min([(norm(sample - center), idx) for idx, center in enumerate(centers)])
    return idx

km_predict_two([0.5, 0.5])

7

What do we get if we analyze the second method?

module_frontier(km_predict_two)

[]

This is a problem! We’d expect that norm would be a captured variable in the body of km_predict_two (and thus that numpy.linalg would be listed in its module frontier), but that isn’t the case. We can inspect the closure variables:

inspect.getclosurevars(km_predict_two)

ClosureVars(nonlocals={}, globals={'centers': array([[ 0.15441674,  0.15065163],
       [ 0.47375581,  0.78146907],
       [ 0.83512659,  0.19018115],
       [ 0.16262154,  0.86710792],
       [ 0.83007508,  0.83832402],
       [ 0.16133578,  0.49974156],
       [ 0.49490377,  0.22475294],
       [ 0.75499895,  0.51576093]])}, builtins={'min': <built-in function min>, 'enumerate': <class 'enumerate'>}, unbound=set())

We can see the cluster centers as well as the min function and the enumerate type. But norm isn’t in the list. Let’s dive deeper. We can use the dis module (and some functionality that was introduced in Python 3.4) to inspect the Python bytecode for a given function:

from dis import Bytecode
for inst in Bytecode(km_predict_two):
    print("%d: %s(%s)" % (inst.offset, inst.opname, inst.argrepr))

0: LOAD_GLOBAL(min)
2: LOAD_CLOSURE(sample)
4: BUILD_TUPLE()
6: LOAD_CONST(<code object <listcomp> at 0x116dffd20, file "<ipython-input-8-5a350184a257>", line 5>)
8: LOAD_CONST('km_predict_two.<locals>.<listcomp>')
10: MAKE_FUNCTION()
12: LOAD_GLOBAL(enumerate)
14: LOAD_GLOBAL(centers)
16: CALL_FUNCTION()
18: GET_ITER()
20: CALL_FUNCTION()
22: CALL_FUNCTION()
24: UNPACK_SEQUENCE()
26: STORE_FAST(_)
28: STORE_FAST(idx)
30: LOAD_FAST(idx)
32: RETURN_VALUE()

Ah ha! The body of our list comprehension, which contains the call to norm, is a separate code object that has been stored in a constant. Let’s look at the constants for our function:

km_predict_two.__code__.co_consts

(None,
 <code object <listcomp> at 0x116dffd20, file "<ipython-input-8-5a350184a257>", line 5>,
 'km_predict_two.<locals>.<listcomp>')

We can see the code object in the constant list and use dis to disassemble it as well:

for inst in Bytecode(km_predict_two.__code__.co_consts[1]):
    print("%d: %s(%s)" % (inst.offset, inst.opname, inst.argrepr))

0: BUILD_LIST()
2: LOAD_FAST(.0)
4: FOR_ITER(to 30)
6: UNPACK_SEQUENCE()
8: STORE_FAST(idx)
10: STORE_FAST(center)
12: LOAD_GLOBAL(norm)
14: LOAD_DEREF(sample)
16: LOAD_FAST(center)
18: BINARY_SUBTRACT()
20: CALL_FUNCTION()
22: LOAD_FAST(idx)
24: BUILD_TUPLE()
26: LIST_APPEND()
28: JUMP_ABSOLUTE()
30: RETURN_VALUE()

Once we’ve done so, we can see that the list comprehension has loaded norm from a global, which we can then resolve and inspect:

km_predict_two.__globals__["norm"]

<function numpy.linalg.linalg.norm>

_.__module__

'numpy.linalg.linalg'

Nested functions and lambda expressions

We can see a similar problem if we look at a function with local definitions (note that there is no need for the nesting in this example other than to expose a limitation of our technique):

import sys

def km_predict_three(sample):
    # unnecessary nested function
    def find_best(sample):
        (n, i) = (sys.float_info.max, -1)
        for idx, center in enumerate(centers):
            (n, i) = min((n, i), (norm(sample - center), idx))
        return i
    return find_best(sample)

km_predict_three([0.5, 0.5])

7

module_frontier(km_predict_three)

[]

In this case, we can see that Python compiles these nested functions in essentially the same way it compiles the bodies of list comprehensions:

from dis import Bytecode
for inst in Bytecode(km_predict_three):
    print("%d: %s(%s)" % (inst.offset, inst.opname, inst.argrepr))

0: LOAD_CONST(<code object find_best at 0x116e0a390, file "<ipython-input-17-e19a1ac37885>", line 5>)
2: LOAD_CONST('km_predict_three.<locals>.find_best')
4: MAKE_FUNCTION()
6: STORE_FAST(find_best)
8: LOAD_FAST(find_best)
10: LOAD_FAST(sample)
12: CALL_FUNCTION()
14: RETURN_VALUE()

But we can inspect the nested function just as we did the list comprehension. Let’s do that but just look at the load instructions; we’ll see that we’ve loaded sys and norm as we’d expect.

from dis import Bytecode
for inst in [op for op in Bytecode(km_predict_three.__code__.co_consts[1]) if "LOAD_" in op.opname]:
    print("%d: %s(%s)" % (inst.offset, inst.opname, inst.argrepr))

0: LOAD_GLOBAL(sys)
2: LOAD_ATTR(float_info)
4: LOAD_ATTR(max)
6: LOAD_CONST(-1)
16: LOAD_GLOBAL(enumerate)
18: LOAD_GLOBAL(centers)
32: LOAD_GLOBAL(min)
34: LOAD_FAST(n)
36: LOAD_FAST(i)
40: LOAD_GLOBAL(norm)
42: LOAD_FAST(sample)
44: LOAD_FAST(center)
50: LOAD_FAST(idx)
66: LOAD_FAST(i)

Predictably, we can also see a similar problem if we analyze a function with lambda expressions:

def km_predict_four(sample):
    _, idx = min(map(lambda tup: (norm(sample - tup[1]), tup[0]), enumerate(centers)))
    return idx

km_predict_four([0.5, 0.5])

7

module_frontier(km_predict_four)

[]

Explicit imports

Let’s look at what happens when we import modules inside the function we’re analyzing. Because of the semantics of import in Python, the module dependency list for this function will depend on whether or not we’ve already imported numpy and sys in the global namespace. If we have, we’ll get a reasonable module list; if we haven’t, we’ll get an empty module list. (If you’re running this code in a notebook, you can try it out by restarting the kernel, re-executing the cell with the definition of module_frontier, and then executing this cell.)

def km_predict_five(sample):
    import numpy
    import sys
    from numpy.linalg import norm
    from sys.float_info import max as MAX_FLOAT
    
    (n, i) = (MAX_FLOAT, -1)
    for idx, center in enumerate(centers):
        (n, i) = min((n, i), (norm(sample - center), idx))
    return i

module_frontier(km_predict_five)

['numpy.core.numeric',
 'builtins',
 'numpy.core.umath',
 'numpy.linalg.linalg',
 'numpy.core.multiarray',
 'numpy.core.numerictypes',
 'numpy',
 'sys',
 'numpy.core._methods',
 'numpy.core.fromnumeric',
 'numpy.lib.type_check',
 'numpy.linalg._umath_linalg']

from dis import Bytecode
for inst in Bytecode(km_predict_five):
    print("%d: %s(%r)" % (inst.offset, inst.opname, inst.argval))

0: LOAD_CONST(0)
2: LOAD_CONST(None)
4: IMPORT_NAME('numpy')
6: STORE_FAST('numpy')
8: LOAD_CONST(0)
10: LOAD_CONST(None)
12: IMPORT_NAME('sys')
14: STORE_FAST('sys')
16: LOAD_CONST(0)
18: LOAD_CONST(('norm',))
20: IMPORT_NAME('numpy.linalg')
22: IMPORT_FROM('norm')
24: STORE_FAST('norm')
26: POP_TOP(None)
28: LOAD_CONST(0)
30: LOAD_CONST(('max',))
32: IMPORT_NAME('sys.float_info')
34: IMPORT_FROM('max')
36: STORE_FAST('MAX_FLOAT')
38: POP_TOP(None)
40: LOAD_FAST('MAX_FLOAT')
42: LOAD_CONST(-1)
44: ROT_TWO(None)
46: STORE_FAST('n')
48: STORE_FAST('i')
50: SETUP_LOOP(102)
52: LOAD_GLOBAL('enumerate')
54: LOAD_GLOBAL('centers')
56: CALL_FUNCTION(1)
58: GET_ITER(None)
60: FOR_ITER(100)
62: UNPACK_SEQUENCE(2)
64: STORE_FAST('idx')
66: STORE_FAST('center')
68: LOAD_GLOBAL('min')
70: LOAD_FAST('n')
72: LOAD_FAST('i')
74: BUILD_TUPLE(2)
76: LOAD_FAST('norm')
78: LOAD_FAST('sample')
80: LOAD_FAST('center')
82: BINARY_SUBTRACT(None)
84: CALL_FUNCTION(1)
86: LOAD_FAST('idx')
88: BUILD_TUPLE(2)
90: CALL_FUNCTION(2)
92: UNPACK_SEQUENCE(2)
94: STORE_FAST('n')
96: STORE_FAST('i')
98: JUMP_ABSOLUTE(60)
100: POP_BLOCK(None)
102: LOAD_FAST('i')
104: RETURN_VALUE(None)

We can see this more clearly by importing a module in a function’s scope that we haven’t imported into the global namespace:

def example_six():
    import json
    return json.loads("{'this-sure-is': 'confusing'}")

module_frontier(example_six)

[]

inspect.getclosurevars(example_six)

ClosureVars(nonlocals={}, globals={}, builtins={}, unbound={'loads', 'json'})

json is an unbound variable (since it isn’t bound in the enclosing environment of the closure). If it were bound in the global namespace, however, the json we’re referring to in example_six would be captured as a global variable:

import json

module_frontier(example_six)

['json']

inspect.getclosurevars(example_six)

ClosureVars(nonlocals={}, globals={'json': <module 'json' from '/Users/willb/anaconda/lib/python3.6/json/__init__.py'>}, builtins={}, unbound={'loads'})

Obviously, we’d like to return the same module-dependency results for functions that import modules locally independently of whether those modules have been imported into the global namespace. We can look at the bytecode for this function to see what instructions might be relevant:

from dis import Bytecode
for inst in Bytecode(example_six):
    print("%d: %s(%r)" % (inst.offset, inst.opname, inst.argval))

0: LOAD_CONST(0)
2: LOAD_CONST(None)
4: IMPORT_NAME('json')
6: STORE_FAST('json')
8: LOAD_FAST('json')
10: LOAD_ATTR('loads')
12: LOAD_CONST("{'this-sure-is': 'confusing'}")
14: CALL_FUNCTION(1)
16: RETURN_VALUE(None)

Solving problems

To address the cases that the closure-inspecting approach misses, we can inspect the bytecode of each function. (We could also inspect abstract syntax trees, using the ast module, but in general it’s easier to do this sort of work with a lower-level, more regular representation. ASTs have more cases to treat than bytecode.)

Python bytecode is stack-based, meaning that each instruction may take one or more arguments from the stack (in addition to explicit arguments encoded in the instruction). For a more involved analysis, we’d probably want to convert Python bytecode to a representation with explicit operands (like three-address code; see section 2 of Vallée-Rai et al. for a reasonable approach), but let’s see how far we can get by just operating on bytecode.

Identifying interesting bytecodes

We know from the problem cases we examined earlier that we need to worry about a few different kinds of bytecode instructions to find some of the modules that our inspect-based approach missed:

• LOAD_CONST instructions that load code objects (e.g., list comprehension bodies, lambda expressions, or nested functions);
• other LOAD_ instructions that might do the same; and
• IMPORT_NAME instructions that import a module into a function’s namespace.

Let’s extend our technique to also inspect relevant bytecodes. We’ll see how far we can get by just looking at bytecodes in isolation (without modeling the stack or value flow). First, we’ll identify the “interesting” bytecodes and return the modules, functions, or code blocks that they implicate:

def interesting(inst):
    from types import CodeType, FunctionType, ModuleType
    from importlib import import_module
    from functools import reduce
    
    if inst.opname == "IMPORT_NAME":
        path = inst.argval.split(".")
        path[0] = [import_module(path[0])]
        result = reduce(lambda x, a: x + [getattr(x[-1], a)], path)
        return ("modules", result)
    if inst.opname == "LOAD_GLOBAL":
        if inst.argval in globals() and type(globals()[inst.argval]) in [CodeType, FunctionType]:
            return ("code", globals()[inst.argval])
        if inst.argval in globals() and type(globals()[inst.argval]) == ModuleType:
            return ("modules", [globals()[inst.argval]])
        else:
            return None
    if "LOAD_" in inst.opname and type(inst.argval) in [CodeType, FunctionType]:
        return ("code", inst.argval)
    return None

Now we can make a revised version of our module_frontier function. This starts with the same basic approach as the initial function but it also:

• processes the bytecode for each code block transitively referred to in each function,
• processes any modules explicitly imported in code.

def mf_revised(f):
  worklist = [f]
  seen = set()
  mods = set()
  for fn in worklist:
    codeworklist = [fn]
    cvs = inspect.getclosurevars(fn)
    gvars = cvs.globals
    for k, v in gvars.items():
        if inspect.ismodule(v):
            mods.add(v.__name__)
        elif inspect.isfunction(v) and id(v) not in seen:
            seen.add(id(v))
            mods.add(v.__module__)
            worklist.append(v)
        elif hasattr(v, "__module__"):
            mods.add(v.__module__)
    for block in codeworklist:
        for (k, v) in [interesting(inst) for inst in Bytecode(block) if interesting(inst)]:
            if k == "modules":
                newmods = [mod.__name__ for mod in v if hasattr(mod, "__name__")]
                mods.update(set(newmods))
            elif k == "code" and id(v) not in seen:
                seen.add(id(v))
                if hasattr(v, "__module__"):
                    mods.add(v.__module__)
            if(inspect.isfunction(v)):
                worklist.append(v)
            elif(inspect.iscode(v)):
                codeworklist.append(v)
   
  result = list(mods)
  result.sort()
  return result

As you can see, this new approach produces sensible results for all of our examples, including the ones that had confounded the closure-variable approach.

mf_revised(km_predict_one)

['numpy', 'sklearn.cluster.k_means_']

mf_revised(km_predict_two)

['builtins',
 'numpy',
 'numpy.core._methods',
 'numpy.core.fromnumeric',
 'numpy.core.multiarray',
 'numpy.core.numeric',
 'numpy.core.numerictypes',
 'numpy.core.umath',
 'numpy.lib.type_check',
 'numpy.linalg._umath_linalg',
 'numpy.linalg.linalg']

mf_revised(km_predict_three)

['builtins',
 'numpy',
 'numpy.core._methods',
 'numpy.core.fromnumeric',
 'numpy.core.multiarray',
 'numpy.core.numeric',
 'numpy.core.numerictypes',
 'numpy.core.umath',
 'numpy.lib.type_check',
 'numpy.linalg._umath_linalg',
 'numpy.linalg.linalg',
 'sys']

mf_revised(km_predict_four)

['builtins',
 'numpy',
 'numpy.core._methods',
 'numpy.core.fromnumeric',
 'numpy.core.multiarray',
 'numpy.core.numeric',
 'numpy.core.numerictypes',
 'numpy.core.umath',
 'numpy.lib.type_check',
 'numpy.linalg._umath_linalg',
 'numpy.linalg.linalg']

mf_revised(km_predict_five)

['builtins',
 'numpy',
 'numpy.core._methods',
 'numpy.core.fromnumeric',
 'numpy.core.multiarray',
 'numpy.core.numeric',
 'numpy.core.numerictypes',
 'numpy.core.umath',
 'numpy.lib.type_check',
 'numpy.linalg',
 'numpy.linalg._umath_linalg',
 'numpy.linalg.linalg',
 'sys']

This is ongoing work and I hope to cover refinements to and extensions of this technique in future posts; as I mentioned at the beginning of the post, the ultimate goal is a tool to publish functions in notebook cells as self-contained services. It has been fun to learn about the power (and limitations) of the inspect module – as well as a little more about how Python compiles code blocks and nested functions.

Thanks to my friend and colleague Erik Erlandson for suggesting improvements to the presentation of this post.