Best way to handle imports used by only some Python classes

Question

I'm making a Python package, inside a module, I have several python classes, but only one of them uses a specific package (tensorflow), which is installed using the extras_require option in setup.py file, since it's a heavy dependency and it's used in a very small part of the project.

Let's say I have the class MyClassRegular, and MyClassTF in the same module, only the second one needs tensorflow, I was importing the package at the top level of the file using:

try:
    import tensorflow as tf
except ModuleNotFoundError:
    logging.error("Tensorflow not found, pip install tensorflow to use MyClassTF")

So this comes with two problems:

If as a user, I'm importing MyClassRegular, it will make warnings about a package I don't even need or care about because I'm using a functionality that is not related to tensorflow
If for some reason, I've installed tensorflow, it could start making warning messages, like the cuda version is not right, or not found a GPU, etc, that again, is not related to MyClassRegular.

So what comes to mind is to import the package inside the MyClassTF, I know this could go somehow against PEP 8, but I don't see a better way to handle it. So giving it a try to this option, I come to the problem that if I import the module on the init, it's not recognized by the classes methods:

class MyClassTF:
    def __init__(self):
        try:
            import tensorflow as tf
        except ModuleNotFoundError: 
            logging.error("Tensorflow not found, pip install tensorflow to use MyClassTF") 

    def train(self):
        print(tf.__version__) # <--- tensorflow it's not recognized here

    def predict(self):
        print(tf.__version__) # <--- Again, not recognized

I could assign tensorflow to a variable like this, but it doesn't feel right:

class MyClassTF:
    def __init__(self):
        try:
            import tensorflow as tf
            self.tf = tf
        except ModuleNotFoundError: 
            logging.error("Tensorflow not found, pip install tensorflow")

So, what would be the best pythonic way to handle this?

EDIT: both MyClassRegular, and MyClassTF are imported in the top __init__.py file using

__all__ = ["MyClassRegular", "MyClassTF"]

I'd like to keep them together, as they are related to the same set of functionalities — Rodrigo A, Jun 27 '21 at 23:30

score 3 · Answer 1 · answered Jun 27 '21 at 17:07

Another way, if you want to delay issuing a warning until the class is actually used (which I think is what you were trying to get to) is something

try:
    import tensorflow as tf
except ImportError:
    # Allow the ImportError to pass silently and just assign tf to None
    tf = None


class MyTF:
    def __init__(self):
        if tf is None:
            warnings.warn('pip install tensorflow to use this class')

or something along those lines. No need to do the import in the method body itself, or assign the tensorflow module to an instance attribute, which is doable, but pretty unusual. The above pattern is more common.

score 1 · Answer 2 · edited Jun 27 '21 at 21:08

Huumm not the most straightforward thing and your solution looks ok for this. However, I would try putting the class in a separate file that is not imported elsewhere in your package. Use the same try code on the module level, then users will only see that error if they try and import that package. Also like this, your python linters should be happy I think.

Several packages use this to control behavior, like fuzzywuzzy

try:
    from .StringMatcher import StringMatcher as SequenceMatcher
except ImportError:
    if platform.python_implementation() != "PyPy":
        warnings.warn('Using slow pure-python SequenceMatcher. Install python-Levenshtein to remove this warning')
    from difflib import SequenceMatcher

https://github.com/seatgeek/fuzzywuzzy/blob/9e3d2fe0d8c1b195696d5fbcda78c371dd4a6b8f/fuzzywuzzy/fuzz.py#L7

I think it could work, now that I thinkabout it, it still witll make those warnings since the classes are imported in the __ init __.py file — Rodrigo A, Jun 27 '21 at 16:50

score 0 · Answer 3 · answered Jun 27 '21 at 21:30

To prevent overhead of testing tf each time you instantiate MyTF, I would proceed like this:

try:
    import tensorflow as tf
    class MyTF(object):
        ...

except ImportError:
    class MyTF(object):
        def __init__(self, *args, **kwargs):
            raise RuntimeError("tensorflow library not available, "
                               "please install it to enable MyTF functionalities")

Or if MyTF is a long code, for more readability, put every thing depending of tensorflow in a _internal_tf.py module and then:

try:
    from ._internal_tf import MyTF
except ImportError:
    class MyTF(object):
        def __init__(self, *args, **kwargs):
            raise RuntimeError("tensorflow library not available, "
                               "please install it to enable MyTF functionalities")

Best way to handle imports used by only some Python classes

3 Answers3