Import only functions from a python file

Question

I have many Python files (submission1.py, submission2.py, ... , submissionN.py) in the following format,

#submission1.py
def fun():
   print('some fancy function')

fun()

I want to write a tester to test these submissions. (They are actually homeworks that I am grading.). I have a tester for the fun() which is able to test the function itself. However, my problem is, when I import submission.py, it runs the fun() since it calls it at the end of file.

I know that, using if __name__ == "__main__": is the correct way of handling this issue, however, our submissions does not have it since we did not teach it.

So, my question is, is there any way that I can import only fun() from the submission.py files without running the rest of the python file?

Answer 1

maybe if you just want to import the fun () function from submission.py try

from submission import fun

To perform the function of fun, you must include the fun module

submission.fun()

or if you want to make it easier when calling the fun () function, give it a try from submission import fun as FUN

FUN ()

Answer 2

For simple scripts with just functions the following will work:

submission1.py:

def fun(x):
   print(x)

fun("foo")


def fun2(x):
   print(x)


fun2("bar")

print("debug print")

You can remove all bar the FunctionDef nodes then recompile:

import ast
import types

with open("submission1.py") as f:
   p = ast.parse(f.read())

for node in p.body[:]:
    if not isinstance(node, ast.FunctionDef):
        p.body.remove(node)



module = types.ModuleType("mod")
code = compile(p, "mod.py", 'exec')
sys.modules["mod"] = module
exec(code,  module.__dict__)

import mod

mod.fun("calling fun")
mod.fun2("calling fun2")

Output:

calling fun
calling fun2

The module body contains two Expr and one Print node which we remove in the loop keeping just the FunctionDef's.

[<_ast.FunctionDef object at 0x7fa33357f610>, <_ast.Expr object at 0x7fa330298a90>, 
<_ast.FunctionDef object at 0x7fa330298b90>, <_ast.Expr object at 0x7fa330298cd0>,
 <_ast.Print object at 0x7fa330298dd0>]

So after the loop out body only contains the functions:

[<_ast.FunctionDef object at 0x7f49a786a610>, <_ast.FunctionDef object at 0x7f49a4583b90>]

This will also catch where the functions are called with print which if the student was calling the function from an IDE where the functions have return statements is pretty likely, also to keep any imports of there are any you can keep ast.Import's and ast.ImportFrom's:

submission.py:

from math import *
import datetime

def fun(x):
    print(x)


fun("foo")


def fun2(x):
    return x

def get_date():
    print(pi)
    return datetime.datetime.now()
fun2("bar")

print("debug print")

print(fun2("hello world"))

print(get_date())

Compile then import:

for node in p.body[:]:
    if not isinstance(node, (ast.FunctionDef,ast.Import, ast.ImportFrom)):
        p.body.remove(node)
.....

import mod

mod.fun("calling fun")
print(mod.fun2("calling fun2"))
print(mod.get_date())

Output:

calling fun
calling fun2
3.14159265359
2015-05-09 12:29:02.472329

Lastly if you have some variables declared that you need to use you can keep them using ast.Assign:

submission.py:

from math import *
import datetime

AREA = 25
WIDTH = 35

def fun(x):
    print(x)


fun("foo")


def fun2(x):
    return x

def get_date():
    print(pi)
    return datetime.datetime.now()
fun2("bar")

print("debug print")

print(fun2("hello world"))

print(get_date()

Add ast.Assign:

for node in p.body[:]:
    if not isinstance(node, (ast.FunctionDef,
        ast.Import, ast.ImportFrom,ast.Assign)):
        p.body.remove(node)
....

Output:

calling fun
calling fun2
3.14159265359
2015-05-09 12:34:18.015799
25
35

So it really all depends on how your modules are structured and what they should contain as to what you remove. If there are literally only functions then the first example will do what you want. If there are other parts that need to be kept it is just a matter of adding them to the isinstance check.

The listing of all the abstract grammar definitions is in the cpython source under Parser/Python.asdl.

Answer 3

You could use sys.settrace() to catch function definitions.

Whenever your fun() is defined, you save it somewhere, and you place a stub into the module you are importing, so that it won't get executed.

Assuming that fun() gets defined only once, this code should do the trick:

import sys

fun = None

def stub(*args, **kwargs):
    pass

def wait_for_fun(frame, event, arg):
    global fun

    if frame.f_code.co_filename == '/path/to/module.py':
        if 'fun' in frame.f_globals:
            # The function has just been defined. Save it.
            fun = frame.f_globals['fun']
            # And replace it with our stub.
            frame.f_globals['fun'] = stub

            # Stop tracing the module execution.
            return None

    return wait_for_fun

sys.settrace(wait_for_fun)
import my_module

# Now fun() is available and we can test it.
fun(1, 2, 3)
# We can also put it again inside the module.
# This is important if other functions in the module need it.
my_module.fun = fun

This code can be improved in many ways, but it does its job.