how can I tell if a function can be used in a context manager?

user_15068079 0
Viewed 647

To change the printing precision of numpy array x, I've been using this:

with np.printoptions(precision=2, suppress=True):
    print(x)

I want to do the same for torch tensor aTensor, but the below doesn't work (I get AttributeError: __enter__ ):

with torch.set_printoptions(precision=2):
    print(aTensor)

I'm new to python and looked around and read that for a function to be used in a context manager, it needs to have __enter__ and __exit__. But when I tried to check, I get that neither np.printoptions nor torch.set_printoptions has __enter__ as an attribute: hasattr(np.printoptions, "__enter__") returns False and the same for hasattr(torch.set_printoptions, "__enter__").

But the former can be used in the context manager but the latter not. Why is this?

The direct problem of torch tensors print precision I can deal with by changing precision, printing and then changing precision back. I'm more interested in learning about the basics of the context manager. Thanks in advance.

torch python
1 Answers

The reason np.printoptions doesn't have an __enter__ attribute is that it's a function that returns a context manager; it's not a context manager itself.

>>> from contextlib import AbstractContextManager
>>> cm = np.printoptions(precision=2, suppress=True)
>>> cm
<contextlib._GeneratorContextManager object at 0x636f6e747874>
>>> isinstance(cm, AbstractContextManager)
True

Note that not all context managers are going to be contextlib._GeneratorContextManager objects; numpy just happens to be using contextlib from the standard library to create this context manager.

To answer your literal question, that last line of code will check if something is a context manager; you can check if it's an instance of contextlib.AbstractContextManager. That's how you should do it if you need to check if something is a context manager in your code. If you just need to check adhoc for your own knowledge and can't be bothered to do it that way, then you can: check in the REPL that it has __enter__ and __exit__ attributes, either by trying to auto-complete them or by using dir(); try to use it as a context manager; or check the documentation/implementation.

Unlike np.printoptions, torch.set_printoptions doesn't even return a context manager, which is why you got that AttributeError. However, you can create your own context manager to handle torch.set_printoptions for you, which you could then use in the same way as np.printoptions. Here's an example; I haven't tested this, but any potential issues can be resolved. You can see relevant code here.

import contextlib
import copy

import torch

@contextlib.contextmanager
def torch_set_printoptions_cm(*args, **kwargs):
    try:
        # be warned, torch._tensor_str is a private module,
        # not bound by API guarantees
        original_options = torch._tensor_str.PRINT_OPTS
        torch._tensor_str.PRINT_OPTS = copy.copy(original_options)

        torch.set_printoptions(*args, **kwargs)
        yield torch._tensor_str.PRINT_OPTS
    finally:
        torch._tensor_str.PRINT_OPTS = original_options

You could then do what you tried:

with torch_set_printoptions_cm(precision=2):
    print(aTensor)
user_8923830
Related