Tuesday, July 08, 2008

Simple stuff with import hooks in Python

Yesterday Jim an I had a tutorial session at EuroPython about dynamic compilation in Python. We brought up the topics of import hooks (PEP 302) since we have successfully used them as an opportunity to create code dynamically. The code example we demonstrated for that was one of the actual import hooks that we had used in a Jython setup. Even if it was no more than 80 lines of code, it might not have been the most accessible example. A few people asked me afterwords if I had a more simple example, so by public request, here is a simple meta_path hook that prevents the user from importing some modules.
import sys

class Restriction(object):
__forbidden = set()
@classmethod
def add(cls, module_name):
cls.__forbidden.add(module_name)
def find_module(self, module_name, package_path):
if package_path: return
if module_name in self.__forbidden:
return self
def load_module(self, module_name):
raise ImportError("Restricted")

sys.meta_path.append(Restriction())
add = Restriction.add
del Restriction
If we walk through this class from the top we first of all have a set containing the names of the modules that we will not allow the user to import and a method for adding modules to that set.
The next method, find_module, is the method invoked by the import subsystem. It is responsible for locating the requested module and return an object capable of loading it. The arguments passed to find_module are the fully qualified module name and, if the module is a sub module of a package, the path where that package was found. If the module cannot be found one should either return None or raise an ImportError. This is a handy way to delegate to the next import hook or the default import behavior. If it returns a loader object on the other hand, no other import mechanism will be attempted, which of course is useful in this case. So in this implementation we return self if the module name is found to be the one of the modules that we want to prevent the user from importing.
The loader object should implement the method load_module, which takes only one argument, the fully qualified module name and returns the corresponding module or raises an ImportError if the import fails. In this case we know that load_module is only ever invoked if we want the import to fail, therefore we always raise an ImportError.
There really isn't much more to it. It should be noted however that this isn't a good way to implement security restrictions in Python, since it is possible for any user code to remove the import hook from sys.meta_path, but I still think it makes for a good introductory example to import hooks.

Happy hacking!

No comments: