Executing modules specified by strings

Question

I am developing a Django backend for an online course platform. The backend runs the code submitted by the student. Here is a working example for running a student code consisting of three modules:

import importlib
import sys

util_a = """
def foo_a():
    print('Yes!')
"""

util_b = """
from util_a import foo_a
def foo_b():
    foo_a()
"""

main = """
from util_b import foo_b
foo_b()
"""

def process_module(code, name):
    module = importlib.util.module_from_spec(
        importlib.util.spec_from_loader(name, loader=None))
    compiled_code = compile(code, '<string>', 'exec')
    exec(compiled_code, module.__dict__)
    sys.modules[name] = module

for code, name in [(util_a, 'util_a'), (util_b, 'util_b'), (main, 'main')]:
    process_module(code, name)

The problem is that the util modules have to be specified in the correct order (i.e. [(util_b, 'util_b'), (util_a, 'util_a'), (main, 'main')] would not work), whereas I want to be able to only specify which module is main, and the rest should happen automatically, just like it would had the modules been real files.

So, how can I modify this code to make it work with util modules specified in any order?

P.S. I have a complicated solution using ast to get the list of modules imported by each module, and topological sort to execute the modules in the correct order. I am looking for a simpler way.

Answer 1

Here is the complicated solution. The original version was generated using ChatGPT. However, I have improved and re-factored it enough to be comfortable with SO policies. Comments appreciated.

import importlib
import sys
import ast
from collections import defaultdict, deque

util_a = """
def foo_a():
    print('Yes!')
"""

util_b = """
from util_a import foo_a
def foo_b():
    foo_a()
"""

main = """
import math

from util_b import foo_b
foo_b()
"""

def get_correct_order(modules):
    def get_imported_modules(code):
        imported_modules = []
        tree = ast.parse(code)
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    imported_modules.append(alias.name)
            elif isinstance(node, ast.ImportFrom):
                imported_modules.append(node.module)
        return filter(
            lambda name:
                name not in sys.builtin_module_names and
                name not in sys.modules,
            imported_modules)

    # Build dependency graph
    graph = defaultdict(set)
    for name, code in modules:
        imported_modules = get_imported_modules(code)
        graph[name].update(imported_modules)

    # Perform topological sorting
    sorted_modules = []
    visited = set()
    queue = deque()

    # Add modules without dependencies to the queue
    for module in graph:
        if not graph[module]:
            queue.append(module)

    while queue:
        module = queue.popleft()
        sorted_modules.append(module)
        visited.add(module)

        # Update dependencies of other modules
        for dependent_module, dependencies in graph.items():
            if module in dependencies:
                dependencies.remove(module)
                if not dependencies and dependent_module not in visited:
                    queue.append(dependent_module)

    return sorted_modules


def process_module(name, code):
    module = importlib.util.module_from_spec(
        importlib.util.spec_from_loader(name, loader=None))
    compiled_code = compile(code, '<string>', 'exec')
    exec(compiled_code, module.__dict__)
    sys.modules[name] = module


modules = [('main', main), ('util_a', util_a), ('util_b', util_b)]
modules_dict = {name: (name, code) for name, code in modules}

# Process modules in the correct order
for name in get_correct_order(modules):
    name, code = modules_dict[name]
    process_module(name, code)