How to properly import a function for use in timeit

Question

I'm using the timeit function to run some metrics on two methods I've written, and I'm encountering an issue. See the simplified example files below:

fileA.py, this runs stand alone, and I'm trying to test the function, foo with timeit:

if (len(sys.argv) < 2):
    print "Need a command line argument, exiting"
    sys.exit(1)

def foo(n):
    #does some stuff, returns an int

This function works fine standalone, I can run it as ./fileA.py 5

I have a similar file, fileB.py that is a different implementation of foo, but is exactly the same except the inside of foo.

Now I have a fileC, where I'm trying to use timeit on both fileA and fileB(fileB part omitted):

for n in range(0,10):
    setupStr = 'from fileA import foo'
    setupStr += '; from __main__ import n'
    mytime = timeit.Timer('foo(n)', setupStr)
    timeTaken = mytime.timeit(1)
    print #results, not important

The issue that I am getting, is that when I run fileC.py, I get "Need a command line argument, exitng" and it exits. I understand that this is coming from fileA.py, but why is this portion of code being run at all? I was under the impression that the way timeit works, I am only importing foo from fileA. Am I assuming wrong? If so, what would be the best way to rectify this, so that I can test the method foo(n) and keep fileA's ability to run standalone with the required command line argument?

Answer 1

Python has to evaluate the whole fileA module to resolve from fileA import foo - it would be perfectly legitimate to have foo be a module-level variable name rather than a function, created by code that would have to be evaluated to get the correct value. To avoid this, wrap the code that's supposed to only be run when fileA is being called standalone in:

if __name__ == '__main__':
    if (len(sys.argv) < 2):
        ...

See the accepted answer to this question for a more detailed explanation: What does if __name__ == "__main__": do?