Python: intercept a class loading action

Question

Summary: when a certain python module is imported, I want to be able to intercept this action, and instead of loading the required class, I want to load another class of my choice.

Reason: I am working on some legacy code. I need to write some unit test code before I start some enhancement/refactoring. The code imports a certain module which will fail in a unit test setting, however. (Because of database server dependency)

Pseduo Code:

from LegacyDataLoader import load_me_data
...
def do_something():
   data = load_me_data()

So, ideally, when python excutes the import line above in a unit test, an alternative class, says MockDataLoader, is loaded instead.

I am still using 2.4.3. I suppose there is an import hook I can manipulate

Edit

Thanks a lot for the answers so far. They are all very helpful.

One particular type of suggestion is about manipulation of PYTHONPATH. It does not work in my case. So I will elaborate my particular situation here.

The original codebase is organised in this way

./dir1/myapp/database/LegacyDataLoader.py
./dir1/myapp/database/Other.py
./dir1/myapp/database/__init__.py
./dir1/myapp/__init__.py

My goal is to enhance the Other class in the Other module. But since it is legacy code, I do not feel comfortable working on it without strapping a test suite around it first.

Now I introduce this unit test code

./unit_test/test.py

The content is simply:

from myapp.database.Other import Other

def test1():
            o = Other()
            o.do_something()

if __name__ == "__main__":
            test1()

When the CI server runs the above test, the test fails. It is because class Other uses LegacyDataLoader, and LegacydataLoader cannot establish database connection to the db server from the CI box.

Now let's add a fake class as suggested:

./unit_test_fake/myapp/database/LegacyDataLoader.py
./unit_test_fake/myapp/database/__init__.py
./unit_test_fake/myapp/__init__.py

Modify the PYTHONPATH to

export PYTHONPATH=unit_test_fake:dir1:unit_test

Now the test fails for another reason

  File "unit_test/test.py", line 1, in <module>
    from myapp.database.Other import Other
ImportError: No module named Other

It has something to do with the way python resolves classes/attributes in a module

Answer 1

Well, if the import fails by raising an exception, you could put it in a try...except loop:

try:
    from LegacyDataLoader import load_me_data
except: # put error that occurs here, so as not to mask actual problems
    from MockDataLoader import load_me_data

Is that what you're looking for? If it fails, but doesn't raise an exception, you could have it run the unit test with a special command line tag, like --unittest, like this:

import sys
if "--unittest" in sys.argv:
    from MockDataLoader import load_me_data
else:
    from LegacyDataLoader import load_me_data

Answer 2

There are cleaner ways to do this, but I'll assume that you can't modify the file containing from LegacyDataLoader import load_me_data.

The simplest thing to do is probably to create a new directory called testing_shims, and create LegacyDataLoader.py file in it. In that file, define whatever fake load_me_data you like. When running the unit tests, put testing_shims into your PYTHONPATH environment variable as the first directory. Alternately, you can modify your test runner to insert testing_shims as the first value in sys.path.

This way, your file will be found when importing LegacyDataLoader, and your code will be loaded instead of the real code.

Answer 3

The import statement just grabs stuff from sys.modules if a matching name is found there, so the simplest thing is to make sure you insert your own module into sys.modules under the target name before anything else tries to import the real thing.

# in test code
import sys
import MockDataLoader
sys.modules['LegacyDataLoader'] = MockDataLoader

import module_under_test

There are a handful of variations on the theme, but that basic approach should work fine to do what you describe in the question. A slightly simpler approach would be this, using just a mock function to replace the one in question:

# in test code
import module_under_test

def mock_load_me_data():
    # do mock stuff here

module_under_test.load_me_data = mock_load_me_data

That simply replaces the appropriate name right in the module itself, so when you invoke the code under test, presumably do_something() in your question, it calls your mock routine.

Answer 4

You can intercept import and from ... import statements by defining your own __import__ function and assigning it to __builtin__.__import__ (make sure to save the previous value, since your override will no doubt want to delegate to it; and you'll need to import __builtin__ to get the builtin-objects module).

For example (Py2.4 specific, since that's what you're asking about), save in aim.py the following:

import __builtin__
realimp = __builtin__.__import__
def my_import(name, globals={}, locals={}, fromlist=[]):
  print 'importing', name, fromlist
  return realimp(name, globals, locals, fromlist)
__builtin__.__import__ = my_import

from os import path

and now:

$ python2.4 aim.py
importing os ('path',)

So this lets you intercept any specific import request you want, and alter the imported module[s] as you wish before you return them -- see the specs here. This is the kind of "hook" you're looking for, right?