Import optional packages in a nicer way

[akihironitta]

🚀 Feature

I would like to suggest to implement a context manager class in order to handle optional imports better. With the context manager implemented, importing optional packages and deferring raising ModuleNotFoundError will look like:

with ContextManager() as cm:
    import torchvision  # doesn't raise `ModuleNotFoundError` here even if it fails.

class SomeClass:
    def __init__(self, ...):
        cm.check()  # raises `ModuleNotFoundError` if `import torchvision` failed.
        ...

Motivation

The suggested context manager reduces duplicate information in the codebase. Currently, there are redundant lines such as try: ... except ModuleNotFoundErorr: ..., _PACKAGE_AVAILABLE and raise ModuleNotFoundError(...).

Example

https://github.com/PyTorchLightning/pytorch-lightning-bolts/blob/4d254fde6112b21436003028d553a726bf7ea6ef/pl_bolts/datamodules/cifar10_datamodule.py#L12-L20 https://github.com/PyTorchLightning/pytorch-lightning-bolts/blob/4d254fde6112b21436003028d553a726bf7ea6ef/pl_bolts/datamodules/cifar10_datamodule.py#L86-L89

Pitch

Here is a minimal implementation of context manager and its behaviour.

Example

class ContextManager:
    def __init__(self):
        self.exc = None
        
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.exc = exc_type, exc_value, traceback
        return True  # Ignores exception if True and reraises exception if False

    def check(self):
        if self.exc is not None:
            exc_type, exc_value, traceback = self.exc
            raise exc_type(exc_value).with_traceback(traceback)

# Here's how it behaves
print("Try to import.")
with ContextManager() as cm:
    import aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa

print("Tried to import.")
cm.check()

Try to import.
Tried to import.
Traceback (most recent call last):
  File "main.py", line 22, in <module>
    cm.check()
  File "main.py", line 15, in check
    raise exc_type(exc_value).with_traceback(traceback)
  File "main.py", line 19, in <module>
    import aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
ModuleNotFoundError: No module named 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa'

Alternatives

As suggested in https://github.com/PyTorchLightning/pytorch-lightning-bolts/pull/338#discussion_r518756049 by @Borda, we could define global variables like _TORCHVISION_AVAILABLE in pl_bolts/__init__.py, but in that way, we will have to add new variables like _NEW_PACKAGE_AVAILABLE as we add new optional packages to Bolts.

See also

Built-in Types - Context Manager Types
Stack Overflow - How to safely handle an exception inside a context manager

Other comments

We can implement the suggested context manager in PL and use it from Bolts. If it should be implemented in PL, could you transfer this issue to PL repo? GitHub Docs - Transferring an issue to another repository

Let me know if it sounds reasonable or too much engineering!

[Borda]

I would agree with the first usage example, define _TORCHVISION_AVAILABLE in relevant upper level init or even in pl_bolts.__init__ but then we do not need to define extra Contex class, then the limiting importing is handled inside https://github.com/PyTorchLightning/pytorch-lightning-bolts/blob/f00e85dd204fa2434c47d9ae20614b0a2baeb1f0/pl_bolts/utils/warnings.py#L7 but looking at it again, it looks cute: @ananyahjha93 what do you think?

[akihironitta]

Found some related PRs in other repos, so let me leave them here:

• https://github.com/optuna/optuna/pull/1315 <- This is pretty much the same as what I am suggesting here.
• https://github.com/tensorflow/datasets/pull/1947

[akihironitta]

@ananyahjha93 Any thoughts on this? :]

[akihironitta]

@Borda @ananyahjha93 ~Can I give it a go for this? otherwise,~ I'll restart to work on #338 defining global variables like _TORCHVISION_AVAILABLE in pl_bolts/__init__.py for now without implementing the context manager suggested here as we may need some more time to discuss this.

[Borda]

I start playing with conditional imports at https://github.com/PyTorchLightning/pytorch-lightning/pull/4859 https://github.com/PyTorchLightning/pytorch-lightning/pull/4860 but it seems it does not work as I would expect, happy to hear some suggestions :]