List all classes in a package directory

Question

I have a package directory pkg with several classes that I would like to build into a convenient dict property.

The structure of pkg/ looks like:

pkg/base.py:

class _MyBase(object):
    pass

pkg/foo.py:

from .base import _MyBase

class Foo(_MyBase):
     pass

And in pkg/__init__.py, it is a bit clunky, but once pkg is imported, a all_my_base_classes dict is built with a key of the class name, and value of the class object. The classes are all subclasses of pkg.base._MyBase.

import os
import sys
import pkgutil

import base
# I don't want to import foo, bar, or whatever other file is in pkg/

all_my_base_classes = {}
pkg_dir = os.path.dirname(__file__)
for (module_loader, name, ispkg) in pkgutil.iter_modules([pkg_dir]):
    exec('import ' + name)
    pkg_name = __name__ + '.' + name
    obj = sys.modules[pkg_name]
    for dir_name in dir(obj):
        if dir_name.startswith('_'):
            continue
        dir_obj = getattr(obj, dir_name)
        if issubclass(dir_obj, base._MyBase):
            all_my_base_classes[dir_name] = dir_obj

Running it from an interactive Python shell, one directory below pkg/:

>>> import pkg
>>> pkg.all_my_base_classes
{'Foo': <class 'pkg.foo.Foo'>}

So it works as expected, but pkg/__init__.py is pretty terrible looking. How can it be better?

Answer 1

Since the classes are all subclasses of _MyBase, they can be accessed via _MyBase.__subclasses__() after they have been imported:

for (module_loader, name, ispkg) in pkgutil.iter_modules([pkg_dir]):
    importlib.import_module('.' + name, __package__)

all_my_base_classes = {cls.__name__: cls for cls in base._MyBase.__subclasses__()}

For importing the modules, I followed the advice of Nihathrael.

Answer 2

Using exec() is generally regarded as bad practice (https://stackoverflow.com/questions/1933451/why-should-exec-and-eval-be-avoided). Since you are trying to import a module, you can do the same by using the importlib package.

Python 2: https://docs.python.org/2/library/importlib.html

Python 3: https://docs.python.org/3/library/importlib.html#module-importlib

Using it should clean up the rest of the code as well.

You might also be interested in this discussion for more possibilities regarding imports: https://stackoverflow.com/questions/1057431/loading-all-modules-in-a-folder-in-python

Edit 2: Removed my code example, as @janne-karila's solution is even shorter and more concise.

Answer 3

The limitations of using the __subclasses()__ approach are:

1. You need to load the classes (i.e. have them imported), otherwise it will not work.
2. If you have multiple layers of inheritance, you will need to recurse. For example: class A -> class B(A) -> class C(B). A.__subclasses__() will only give you class B not class C. This is because class C inherits from class B not A!

Therefore, to address these 2 limitations I have a slightly different approach in addition to the awesome __subclasses__ answer above.

This function will get all the classes in the package:

def get_classes_from_package_recursively(package: str) -> list[type]:
    """Return a list of classes inside a given package (recurse thorugh any sub-packages).

    Keyword arguments:
    package -- package represented as a string. Must not be relative.
    """
    classes_in_package = []
    # Go through the modules in the package
    for _importer, module_name, is_package in pkgutil.iter_modules(importlib.import_module(package).__path__):
        full_module_name = f"{package}.{module_name}"
        # Recurse through any sub-packages
        if is_package:
            classes_in_subpackage = get_classes_from_package_recursively(package=full_module_name)
            classes_in_package.extend(classes_in_subpackage)

        # Load the module for inspection
        module = importlib.import_module(full_module_name)

        # Iterate through all the objects in the module and
        # using the lambda, filter for class objects and only objects that exist within the module
        for _name, obj in inspect.getmembers(
            module,
            lambda member, module_name=full_module_name: inspect.isclass(member) and member.__module__ == module_name,
        ):
            classes_in_package.append(obj)
    return classes_in_package

Example of how to use the function above:

my_classes = get_classes_from_package_recursively(package="src.database.models")

my_classes will now look like this [MyModel1, MyModel2, MyModel3, etc...]

After getting this, simply filter through the list like this:

my_subclasses = [class_ for class_ in my_classes if issubclass(class_, parent_class) and class_ is not parent_class]