Making Python run a few lines before my script

问题:

I need to run a script foo.py, but I need to also insert some debugging lines to run before the code in foo.py. Currently I just put those lines in foo.py and I'm careful not to commit that to Git, but I don't like this solution.

What I want is a separate file bar.py that I don't commit to Git. Then I want to run:

python /somewhere/bar.py /somewhere_else/foo.py

What I want this to do is first run some lines of code in bar.py, and then run foo.py as __main__. It should be in the same process that the bar.py lines ran in, otherwise the debugging lines won't help.

Is there a way to make bar.py do this?

Someone suggested this:

import imp
import sys

# Debugging code here

fp, pathname, description = imp.find_module(sys.argv[1])
imp.load_module('__main__', fp, pathname, description)

The problem with that is that because it uses import machinery, I need to be on the same folder as foo.py to run that. I don't want that. I want to simply put in the full path to foo.py.

Also: The solution needs to work with .pyc files as well.

回答1:

Python has a mechanism for running code at startup; the site module.

"This module is automatically imported during initialization."

The site module will attempt to import a module named sitecustomize before __main__ is imported. It will also attempt to import a module named usercustomize if your environment instructs it to.

For example, you could put a sitecustomize.py file in your site-packages folder that contains this:

import imp

import os

if 'MY_STARTUP_FILE' in os.environ:
    try:
        file_path = os.environ['MY_STARTUP_FILE']
        folder, file_name = os.path.split(file_path)
        module_name, _ = os.path.splitext(file_name)
        fp, pathname, description = imp.find_module(module_name, [folder])
    except Exception as e:
        # Broad exception handling since sitecustomize exceptions are ignored
        print "There was a problem finding startup file", file_path
        print repr(e)
        exit()

    try:
        imp.load_module(module_name, fp, pathname, description)
    except Exception as e:
        print "There was a problem loading startup file: ", file_path
        print repr(e)
        exit()
    finally:
        # "the caller is responsible for closing the file argument" from imp docs
        if fp:
            fp.close()

Then you could run your script like this:

MY_STARTUP_FILE=/somewhere/bar.py python /somewhere_else/foo.py

• You could run any script before foo.py without needing to add code to reimport __main__.
• Run export MY_STARTUP_FILE=/somewhere/bar.py and not need to reference it every time

回答2:

You can use execfile() if the file is .py and uncompyle2 if the file is .pyc.

Let's say you have your file structure like:

test|-- foo.py
    |-- bar
         |--bar.py   

foo.py

import sys

a = 1
print ('debugging...')

# run the other file
if sys.argv[1].endswith('.py'): # if .py run right away
    execfile(sys.argv[1], globals(), locals())
elif sys.argv[1].endswith('.pyc'): # if .pyc, first uncompyle, then run
    import uncompyle2
    from StringIO import StringIO
    f = StringIO()
    uncompyle2.uncompyle_file(sys.argv[1], f)
    f.seek(0)
    exec(f.read(), globals(), locals())

bar.py

print a
print 'real job'

And in test/, if you do:

$ python foo.py bar/bar.py
$ python foo.py bar/bar.pyc

Both, outputs the same:

debugging...
1
real job

Please also see this answer.

回答3:

You probably have something along the lines of:

if __name__ == '__main__':
    # some code

Instead, write your code in a function main() in foo and then do:

if __name__ == '__main__':
    main()

Then, in bar, you can import foo and call foo.main().

Additionaly, if you need to change the working directory, you can use the os.chdir(path) method, e.g. os.chdir('path/of/bar') .

回答4:

bar.py would have to behave like the Python interpreter itself and run foo.py (or foo.pyc, as you asked for that) as if it was the main script. This is surprisingly hard. My 90% solution to do that looks like so:

def run_script_as_main(cmdline):
    # It is crucial to import locally what we need as we
    # later remove everything from __main__
    import sys, imp, traceback, os
    # Patch sys.argv
    sys.argv = cmdline
    # Clear the __main__ namespace and set it up to run
    # the secondary program
    maindict = sys.modules["__main__"].__dict__
    builtins = maindict['__builtins__']
    maindict.clear()
    maindict['__file__'] = cmdline[0]
    maindict['__builtins__'] = builtins
    maindict['__name__'] = "__main__"
    maindict['__doc__'] = None
    # Python prepends a script's location to sys.path
    sys.path[0] = os.path.dirname(os.path.abspath(cmdline[0]))
    # Treat everything as a Python source file, except it
    # ends in '.pyc'
    loader = imp.load_source
    if maindict["__file__"].endswith(".pyc"):
        loader = imp.load_compiled
    with open(cmdline[0], 'rb') as f:
        try:
            loader('__main__', maindict["__file__"], f)
        except Exception:
            # In case of an exception, remove this script from the
            # stack trace; if you don't care seeing some bar.py in
            # traceback, you can leave that out.
            ex_type, ex_value, ex_traceback = sys.exc_info()
            tb = traceback.extract_tb(ex_traceback, None)[1:]
            sys.stderr.write("Traceback (most recent call last):\n")
            for line in traceback.format_list(tb):
                sys.stderr.write(line)
            for line in traceback.format_exception_only(ex_type, ex_value):
                sys.stderr.write(line)

This mimics the default behavior of the Python interpreter relatively closely. It sets system globals __name__, __file__, sys.argv, inserts the script location into sys.path, clears the global namespace etc.

You would copy that code to bar.py or import it from somewhere and use it like so:

if __name__ == "__main__":
    # You debugging setup goes here
    ...
    # Run the Python program given as argv[1]
    run_script_as_main(sys.argv[1:])

Then, given this:

$ find so-foo-bar
so-foo-bar
so-foo-bar/debugaid
so-foo-bar/debugaid/bar.py
so-foo-bar/foo.py

and foo.py looking like this:

import sys

if __name__ == "__main__":
    print "My name is", __name__
    print "My file is", __file__
    print "My command line is", sys.argv
    print "First 2 path items", sys.path[:2]
    print "Globals", globals().keys()
    raise Exception("foo")

... running foo.py via bar.py gives you that:

    $ python so-foo-bar/debugaid/bar.py so-foo-bar/foo.py
    My name is __main__
    My file is so-foo-bar/foo.py
    My command line is ['so-foo-bar/foo.py']
    First 2 path items ['~/so-foo-bar', '/usr/local/...']
    Globals ['__builtins__', '__name__', '__file__', 'sys', '__package__', '__doc__']
    Traceback (most recent call last):
      File "so-foo-bar/foo.py", line 9, in 
        raise Exception("foo")
    Exception: foo

While I'd guess that run_script_as_main is lacking in some interesting details, it's pretty close to the way Python would run foo.py.

回答5:

This solution is what ended up working well for me:

#!/usr/bin/env python

import imp
import sys
import os.path

file_path = sys.argv.pop(1)
assert os.path.exists(file_path)

folder, file_name = os.path.split(file_path)

###############################################################################
#                                                                             #

# Debugging cruft here

#                                                                             #
###############################################################################


module_name, _ = os.path.splitext(file_name)
sys.path.append(folder)
imp.load_module('__main__', *imp.find_module(module_name))

回答6:

The foo.py need not be on the same folder as the main folder. Use

export PYTHONPATH=$HOME/dirWithFoo/:$PYTHONPATH

To add the path that foo resides to Python path. Now you can

from foo import myfunc

myfunc()

And make myfunc do whatever u want

回答7:

Have you tried this?

bar.py

imp.load_source('__main__', '../../foo.py')

for me it runs whatever is in foo.py (before and inside main)

After, maybe what you do in bar.py might be more tricky that my basic test.

The good thing with load_source is that it imports the .pyc if it already exists or initiates a "compile" of the .py and then imports the .pyc.