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

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.