CODEWITHSUNDEEP
python
Robbin
Robbin

Reputation: 304

How do I expose a python class under its containing package without importing the entire package?

I have the following (toy) package structure

root/
 - package1/
   - __init__.py
   - class_a.py
   - class_b.py
 - run.py

In both class_a.py and class_b.py I have a class definition that I want to expose to run.py. If I want to import them this way, I will have to use

from package1.class_a import ClassA  # works but doesn't look nice

I don't like that this shows the class_a.py module, and would rather use the import style

from package1 import ClassA  # what I want

This is also closer to what I see from larger libraries. I found a way to do this by importing the classes in the __init__.py file like so

from class_a import ClassA
from class_b import ClassB

This works fine if it wasn't for one downside: as soon as I import ClassA as I would like (see above), I also immediately 'import' ClassB as, as far as I know, the __init__.py will be run, importing ClassB. In my real scenario, this means I implicitly import a huge class that I use very situationally (which itself imports tensorflow), so I really want to avoid this somehow. Is there a way to create the nice looking imports without automatically importing everything in the package?

Upvotes: 2

Views: 832

Answers (1)

Serge Ballesta
Serge Ballesta

Reputation: 148880

It is possible but require a rather low level customization: you will have to customize the class of your package (possible since Python 3.5). That way, you can declare a __getattr__ member that will be called when you ask for a missing attribute. At that moment, you know that you have to import the relevant module and extract the correct attribute.

The init.py file should contain (names can of course be changed):

import importlib
import sys
import types

class SpecialModule(types.ModuleType):
    """ Customization of a module that is able to dynamically loads submodules.
    
    It is expected to be a plain package (and to be declared in the __init__.py)
    The special attribute is a dictionary attribute name -> relative module name.
    The first time a name is requested, the corresponding module is loaded, and 
    the attribute is binded into the package
    """

    special = {'ClassA': '.class_a', 'ClassB': '.class_b'}

    def __getattr__(self, name):
        if name in self.special:
            m = importlib.import_module(self.special[name], __name__) # import submodule
            o = getattr(m, name)                      # find the required member
            setattr(sys.modules[__name__], name, o)   # bind it into the package
            return o
        else:
            raise AttributeError(f'module {__name__} has no attribute {name}')

sys.modules[__name__].__class__ = SpecialModule       # customize the class of the package

You can now use it that way:

import package1

...
obj = package1.ClassA(...)       # dynamically loads class_a on first call

The downside is that clever IDE that look at the declared member could choke on that and pretend that you are accessing an inexistant member because ClassA is not statically declared in package1/__init__.py. But all will be fine at run time.

As it is a low level customization, it is up to you do know whether it is worth it...

Since 3.7 you could also declare a __gettatr__(name) function directly at the module level.

Upvotes: 2

Related Questions