checkpointing.identifier.func_call.context

  1import inspect
  2from typing import List, Dict, Tuple, Callable, Any, Union
  3from collections import OrderedDict
  4from checkpointing._typing import ReturnValue
  5from types import FrameType
  6import copy
  7
  8
  9class FuncCallContext:
 10    """
 11    Context of information for a function call.
 12    """
 13
 14    def __init__(self, func: Callable[..., ReturnValue], args: Tuple, kwargs: Dict, def_frame: FrameType = None,) -> None:
 15        """
 16        Args:
 17            func: the function object that is being called
 18            args: the non-keywords arguments of the function call
 19            kwargs: the keyword arguments of the function call
 20            def_frame: the frame where the function is defined
 21        """
 22
 23        self.__func: Callable[..., ReturnValue] = func
 24        self.__args: Tuple = args
 25        self.__kwargs: Dict = kwargs
 26        self.__signature = inspect.signature(self.__func)
 27
 28        if def_frame is not None:
 29            self.__locals = copy.copy(def_frame.f_locals)
 30
 31        else:
 32            self.__locals = None
 33
 34    @property
 35    def arguments(self) -> Dict:
 36        """
 37        Dictionary of the function arguments and their actually applied parameter values in the function call.
 38
 39        >>> def foo(a, b=1, c=None, d=4):
 40        ...     pass
 41        >>>
 42        >>> # Equivalent to the context computed for: foo(1, 2, c=3)
 43        >>> ctx = FuncCallContext(foo, (1, 2), {"c": 3})
 44        >>>
 45        >>> dict(ctx.arguments)
 46        {'a': 1, 'b': 2, 'c': 3, 'd': 4}
 47
 48        The varargs and kwargs can also be captured:
 49
 50        >>> def bar(*args, **kwargs):
 51        ...     pass
 52        >>>
 53        >>> ctx = FuncCallContext(bar, (1, 2), {"c": 3})
 54        >>>
 55        >>> ctx.arguments
 56        {'args': (1, 2), 'kwargs': {'c': 3}}
 57        """
 58        args = self.__signature.bind(*self.__args, **self.__kwargs)
 59        args.apply_defaults()
 60        return dict(args.arguments)
 61
 62    @property
 63    def full_name(self) -> str:
 64        """
 65        The full name of the function.
 66        Specifically, the concatenation of a function's module and its [qualified name](https://docs.python.org/3/glossary.html#term-qualified-name).
 67
 68        >>> def foo():
 69        ...     pass
 70        >>>
 71        >>> # Equivalent to the context computed for: foo()
 72        >>> ctx = FuncCallContext(foo, (), {})
 73        >>>
 74        >>> # foo is defined in checkpointing/identifier/func_call/context.py
 75        >>> ctx.full_name
 76        'checkpointing.identifier.func_call.context.foo'
 77        """
 78        return ".".join([self.module, self.qualified_name])
 79
 80    @property
 81    def module(self) -> str:
 82        """
 83        Module where the function is defined.
 84
 85        >>> def foo():
 86        ...     pass
 87        >>>
 88        >>> # Equivalent to the context computed for: foo()
 89        >>> ctx = FuncCallContext(foo, (), {})
 90        >>>
 91        >>> # foo is defined in checkpointing/identifier/func_call/context.py
 92        >>> ctx.module
 93        'checkpointing.identifier.func_call.context'
 94        """
 95
 96        return self.__func.__module__
 97
 98    @property
 99    def name(self) -> str:
100        """
101        Name of the function.
102
103        >>> def foo():
104        ...     pass
105        >>>
106        >>> # Equivalent to the context computed for: foo()
107        >>> ctx = FuncCallContext(foo, (), {})
108        >>> ctx.name
109        'foo'
110        """
111
112        return self.__func.__name__
113
114    @property
115    def qualified_name(self) -> str:
116        """
117        [Qualified name](https://docs.python.org/3/glossary.html#term-qualified-name) of the function.
118
119        >>> class Foo:
120        ...     def bar(self):
121        ...         pass
122        >>>
123        >>> f = Foo()
124        >>>
125        >>> # Equivalent to the context computed for: f.bar()
126        >>> ctx = FuncCallContext(f.bar, (), {})
127        >>> ctx.qualified_name
128        'Foo.bar'
129        """
130
131        return self.__func.__qualname__
132
133    @property
134    def code(self) -> str:
135        r"""
136        The source code of the function, including the function signature,
137        formatted as-is, including any comments.
138
139        >>> def foo(a, b):
140        ...     c = a + b  # add a and b
141        ...     return c
142        >>>
143        >>> # Equivalent to the context computed for: foo(1, 2)
144        >>> ctx = FuncCallContext(foo, (1, 2), {})
145        >>>
146        >>> ctx.code
147        'def foo(a, b):\n    c = a + b  # add a and b\n    return c\n'
148        """
149
150        return inspect.getsource(self.__func)
151
152    def get_nonlocal_variable(self, varname: str) -> Any:
153        r"""
154        Try to get the nonlocal variable `varname` from the function call's
155        `locals()` and `globals()`.
156        If it doesn't exist in neither of them, return `None`.
157        If the caller_frame is not provided, the result is always `None`.
158
159        >>> import inspect
160        >>>
161        >>> a = 1
162        >>> def foo():
163        ...     pass
164        >>>
165        >>> ctx = FuncCallContext(foo, (), {}, inspect.currentframe())
166        >>> ctx.get_nonlocal_variable("a")
167        1
168        >>> ctx.get_nonlocal_variable("b") is None
169        True
170        """
171
172        if self.__locals is not None and varname in self.__locals:
173            return self.__locals[varname]
174
175        if self.__func.__globals__ is not None and varname in self.__func.__globals__:
176            return self.__func.__globals__[varname]
177
178        return None