type_enforced.enforcer

  1from types import (
  2    FunctionType,
  3    MethodType,
  4    GeneratorType,
  5    BuiltinFunctionType,
  6    BuiltinMethodType,
  7    UnionType,
  8)
  9from typing import Type, Union, Sized, Literal, Callable, get_type_hints, Any
 10from functools import update_wrapper, wraps
 11from type_enforced.utils import (
 12    Partial,
 13    GenericConstraint,
 14    DeepMerge,
 15    iterable_types,
 16)
 17
 18
 19class FunctionMethodEnforcer:
 20    def __init__(self, __fn__):
 21        """
 22        Initialize a FunctionMethodEnforcer class object as a wrapper for a passed function `__fn__`.
 23
 24        Requires:
 25
 26            - `__fn__`:
 27                - What: The function to enforce
 28                - Type: function | method | class
 29        """
 30        update_wrapper(self, __fn__)
 31        self.__fn__ = __fn__
 32        self.__outer_self__ = None
 33        # Validate that the passed function or method is a method or function
 34        self.__check_method_function__()
 35        # Get input defaults for the function or method
 36        self.__get_defaults__()
 37
 38    def __get_defaults__(self):
 39        """
 40        Get the default values of the passed function or method and store them in `self.__fn_defaults__`.
 41        """
 42        self.__fn_defaults__ = {}
 43        if self.__fn__.__defaults__ is not None:
 44            # Get the names of all provided default values for args
 45            default_varnames = list(self.__fn__.__code__.co_varnames)[
 46                : self.__fn__.__code__.co_argcount
 47            ][-len(self.__fn__.__defaults__) :]
 48            # Update the output dictionary with the default values
 49            self.__fn_defaults__.update(
 50                dict(zip(default_varnames, self.__fn__.__defaults__))
 51            )
 52        if self.__fn__.__kwdefaults__ is not None:
 53            # Update the output dictionary with the keyword default values
 54            self.__fn_defaults__.update(self.__fn__.__kwdefaults__)
 55
 56    def __get_checkable_types__(self):
 57        """
 58        Creates two class attributes:
 59
 60        - `self.__checkable_types__`:
 61            - What: A dictionary of all annotations as checkable types
 62            - Type: dict
 63
 64        - `self.__return_type__`:
 65            - What: The return type of the function or method
 66            - Type: dict | None
 67        """
 68        if not hasattr(self, "__checkable_types__"):
 69            self.__checkable_types__ = {
 70                key: self.__get_checkable_type__(value)
 71                for key, value in get_type_hints(self.__fn__).items()
 72            }
 73            self.__return_type__ = self.__checkable_types__.pop("return", None)
 74
 75    def __get_checkable_type__(self, annotation):
 76        """
 77        Parses a type annotation and returns a nested dict structure
 78        representing the checkable type(s) for validation.
 79        """
 80
 81        if annotation is None:
 82            return {type(None): None}
 83
 84        # Handle `int | str` syntax (Python 3.10+) and Unions
 85        if (
 86            isinstance(annotation, UnionType)
 87            or getattr(annotation, "__origin__", None) == Union
 88        ):
 89            combined_types = {}
 90            for sub_type in annotation.__args__:
 91                combined_types = DeepMerge(
 92                    combined_types, self.__get_checkable_type__(sub_type)
 93                )
 94            return combined_types
 95
 96        # Handle typing.Literal
 97        if getattr(annotation, "__origin__", None) == Literal:
 98            return {"__extra__": {"__literal__": list(annotation.__args__)}}
 99
100        # Handle generic collections
101        origin = getattr(annotation, "__origin__", None)
102        args = getattr(annotation, "__args__", ())
103
104        if origin == list:
105            if len(args) != 1:
106                raise TypeError(
107                    f"List must have a single type argument, got: {args}"
108                )
109            return {list: self.__get_checkable_type__(args[0])}
110
111        if origin == dict:
112            if len(args) != 2:
113                raise TypeError(
114                    f"Dict must have two type arguments, got: {args}"
115                )
116            key_type = self.__get_checkable_type__(args[0])
117            value_type = self.__get_checkable_type__(args[1])
118            return {dict: (key_type, value_type)}
119
120        if origin == tuple:
121            if len(args) > 2 or len(args) == 1:
122                if Ellipsis in args:
123                    raise TypeError(
124                        "Tuple with Ellipsis must have exactly two type arguments and the second must be Ellipsis."
125                    )
126            if len(args) == 2:
127                if args[0] is Ellipsis:
128                    raise TypeError(
129                        "Tuple with Ellipsis must have exactly two type arguments and the first must not be Ellipsis."
130                    )
131                if args[1] is Ellipsis:
132                    return {tuple: (self.__get_checkable_type__(args[0]), True)}
133            return {
134                tuple: (
135                    tuple(self.__get_checkable_type__(arg) for arg in args),
136                    False,
137                )
138            }
139
140        if origin == set:
141            if len(args) != 1:
142                raise TypeError(
143                    f"Set must have a single type argument, got: {args}"
144                )
145            return {set: self.__get_checkable_type__(args[0])}
146
147        # Handle Sized types
148        if annotation == Sized:
149            return {
150                list: None,
151                tuple: None,
152                dict: None,
153                set: None,
154                str: None,
155                bytes: None,
156                bytearray: None,
157                memoryview: None,
158                range: None,
159            }
160
161        # Handle Callable types
162        if annotation == Callable:
163            return {
164                staticmethod: None,
165                classmethod: None,
166                FunctionType: None,
167                BuiltinFunctionType: None,
168                MethodType: None,
169                BuiltinMethodType: None,
170                GeneratorType: None,
171            }
172
173        if annotation == Any:
174            return {
175                object: None,
176            }
177
178        # Handle Constraints
179        if isinstance(annotation, GenericConstraint):
180            return {"__extra__": {"__constraints__": [annotation]}}
181
182        # Handle standard types
183        if isinstance(annotation, type):
184            return {annotation: None}
185
186        # Hanldle typing.Type (for uninitialized classes)
187        if origin is type and len(args) == 1:
188            return {annotation: None}
189
190        raise TypeError(f"Unsupported type hint: {annotation}")
191
192    def __exception__(self, message):
193        """
194        Usage:
195
196        - Creates a class based exception message
197
198        Requires:
199
200        - `message`:
201            - Type: str
202            - What: The message to warn users with
203        """
204        raise TypeError(
205            f"TypeEnforced Exception ({self.__fn__.__qualname__}): {message}"
206        )
207
208    def __get__(self, obj, objtype):
209        """
210        Overwrite standard __get__ method to return __call__ instead for wrapped class methods.
211
212        Also stores the calling (__get__) `obj` to be passed as an initial argument for `__call__` such that methods can pass `self` correctly.
213        """
214
215        @wraps(self.__fn__)
216        def __get_fn__(*args, **kwargs):
217            return self.__call__(*args, **kwargs)
218
219        self.__outer_self__ = obj
220        return __get_fn__
221
222    def __check_method_function__(self):
223        """
224        Validate that `self.__fn__` is a method or function
225        """
226        if not isinstance(self.__fn__, (MethodType, FunctionType)):
227            raise Exception(
228                f"A non function/method was passed to Enforcer. See the stack trace above for more information."
229            )
230
231    def __call__(self, *args, **kwargs):
232        """
233        This method is used to validate the passed inputs and return the output of the wrapped function or method.
234        """
235        # Special code to pass self as an initial argument
236        # for validation purposes in methods
237        # See: self.__get__
238        if self.__outer_self__ is not None:
239            args = (self.__outer_self__, *args)
240        # Get a dictionary of all annotations as checkable types
241        # Note: This is only done once at first call to avoid redundant calculations
242        self.__get_checkable_types__()
243        # Create a compreshensive dictionary of assigned variables (order matters)
244        assigned_vars = {
245            **self.__fn_defaults__,
246            **dict(zip(self.__fn__.__code__.co_varnames[: len(args)], args)),
247            **kwargs,
248        }
249        # Validate all listed annotations vs the assigned_vars dictionary
250        for key, value in self.__checkable_types__.items():
251            self.__check_type__(assigned_vars.get(key), value, key)
252        # Execute the function callable
253        return_value = self.__fn__(*args, **kwargs)
254        # If a return type was passed, validate the returned object
255        if self.__return_type__ is not None:
256            self.__check_type__(return_value, self.__return_type__, "return")
257        return return_value
258
259    def __check_type__(self, obj, expected, key):
260        """
261        Raises an exception the type of a passed `obj` (parameter) is not in the list of supplied `acceptable_types` for the argument.
262        """
263        # Special case for None
264        if obj is None and type(None) in expected:
265            return
266        extra = expected.get("__extra__", {})
267        expected = {k: v for k, v in expected.items() if k != "__extra__"}
268
269        if isinstance(obj, type):
270            # An uninitialized class is passed, we need to check if the type is in the expected types using Type[obj]
271            obj_type = Type[obj]
272            is_present = obj_type in expected
273        else:
274            obj_type = type(obj)
275            is_present = isinstance(obj, tuple(expected.keys()))
276
277        if not is_present:
278            # Allow for literals to be used to bypass type checks if present
279            literal = extra.get("__literal__", ())
280            if literal:
281                if obj not in literal:
282                    self.__exception__(
283                        f"Type mismatch for typed variable `{key}`. Expected one of the following `{list(expected.keys())}` or a literal value in `{literal}` but got type `{obj_type}` with value `{obj}` instead."
284                    )
285            # Raise an exception if the type is not in the expected types
286            else:
287                self.__exception__(
288                    f"Type mismatch for typed variable `{key}`. Expected one of the following `{list(expected.keys())}` but got `{obj_type}` with value `{obj}` instead."
289                )
290        # If the object_type is in the expected types, we can proceed with validation
291        elif obj_type in iterable_types:
292            subtype = expected.get(obj_type, None)
293            if subtype is None:
294                pass
295            # Recursive validation
296            elif obj_type == list:
297                for idx, item in enumerate(obj):
298                    self.__check_type__(item, subtype, f"{key}[{idx}]")
299            elif obj_type == dict:
300                key_type, val_type = subtype
301                for k, v in obj.items():
302                    self.__check_type__(k, key_type, f"{key}.key[{repr(k)}]")
303                    self.__check_type__(v, val_type, f"{key}[{repr(k)}]")
304            elif obj_type == tuple:
305                expected_args, is_ellipsis = subtype
306                if is_ellipsis:
307                    for idx, item in enumerate(obj):
308                        self.__check_type__(
309                            item, expected_args, f"{key}[{idx}]"
310                        )
311                else:
312                    if len(obj) != len(expected_args):
313                        self.__exception__(
314                            f"Tuple length mismatch for `{key}`. Expected length {len(expected_args)}, got {len(obj)}"
315                        )
316                    for idx, (item, ex) in enumerate(zip(obj, expected_args)):
317                        self.__check_type__(item, ex, f"{key}[{idx}]")
318            elif obj_type == set:
319                for item in obj:
320                    self.__check_type__(item, subtype, f"{key}[{repr(item)}]")
321
322        # Validate constraints if any are present
323        constraints = extra.get("__constraints__", [])
324        for constraint in constraints:
325            constraint_validation_output = constraint.__validate__(key, obj)
326            if constraint_validation_output is not True:
327                self.__exception__(
328                    f"Constraint validation error for variable `{key}` with value `{obj}`. {constraint_validation_output}"
329                )
330
331    def __repr__(self):
332        return f"<type_enforced {self.__fn__.__module__}.{self.__fn__.__qualname__} object at {hex(id(self))}>"
333
334
335def Enforcer(clsFnMethod, enabled):
336    """
337    A wrapper to enforce types within a function or method given argument annotations.
338
339    Each wrapped item is converted into a special `FunctionMethodEnforcer` class object that validates the passed parameters for the function or method when it is called. If a function or method that is passed does not have any annotations, it is not converted into a `FunctionMethodEnforcer` class as no validation is possible.
340
341    If wrapping a class, all methods in the class that meet any of the following criteria will be wrapped individually:
342
343    - Methods with `__call__`
344    - Methods wrapped with `staticmethod` (if python >= 3.10)
345    - Methods wrapped with `classmethod` (if python >= 3.10)
346
347    Requires:
348
349    - `clsFnMethod`:
350        - What: The class, function or method that should have input types enforced
351        - Type: function | method | class
352
353    Optional:
354
355    - `enabled`:
356        - What: A boolean to enable or disable the enforcer
357        - Type: bool
358        - Default: True
359
360    Example Use:
361    ```
362    >>> import type_enforced
363    >>> @type_enforced.Enforcer
364    ... def my_fn(a: int , b: [int, str] =2, c: int =3) -> None:
365    ...     pass
366    ...
367    >>> my_fn(a=1, b=2, c=3)
368    >>> my_fn(a=1, b='2', c=3)
369    >>> my_fn(a='a', b=2, c=3)
370    Traceback (most recent call last):
371      File "<stdin>", line 1, in <module>
372      File "/home/conmak/development/personal/type_enforced/type_enforced/enforcer.py", line 85, in __call__
373        self.__check_type__(assigned_vars.get(key), value, key)
374      File "/home/conmak/development/personal/type_enforced/type_enforced/enforcer.py", line 107, in __check_type__
375        self.__exception__(
376      File "/home/conmak/development/personal/type_enforced/type_enforced/enforcer.py", line 34, in __exception__
377        raise Exception(f"({self.__fn__.__qualname__}): {message}")
378    Exception: (my_fn): Type mismatch for typed variable `a`. Expected one of the following `[<class 'int'>]` but got `<class 'str'>` instead.
379    ```
380    """
381    if not hasattr(clsFnMethod, "__type_enforced_enabled__"):
382        # Special try except clause to handle cases when the object is immutable
383        try:
384            clsFnMethod.__type_enforced_enabled__ = enabled
385        except:
386            return clsFnMethod
387    if not clsFnMethod.__type_enforced_enabled__:
388        return clsFnMethod
389    if isinstance(
390        clsFnMethod, (staticmethod, classmethod, FunctionType, MethodType)
391    ):
392        # Only apply the enforcer if type_hints are present
393        if get_type_hints(clsFnMethod) == {}:
394            return clsFnMethod
395        elif isinstance(clsFnMethod, staticmethod):
396            return staticmethod(FunctionMethodEnforcer(clsFnMethod.__func__))
397        elif isinstance(clsFnMethod, classmethod):
398            return classmethod(FunctionMethodEnforcer(clsFnMethod.__func__))
399        else:
400            return FunctionMethodEnforcer(clsFnMethod)
401    elif hasattr(clsFnMethod, "__dict__"):
402        for key, value in clsFnMethod.__dict__.items():
403            # Skip the __annotate__ method if present in __dict__ as it deletes itself upon invocation
404            # Skip any previously wrapped methods if they are already a FunctionMethodEnforcer
405            if key == "__annotate__" or isinstance(
406                value, FunctionMethodEnforcer
407            ):
408                continue
409            if hasattr(value, "__call__") or isinstance(
410                value, (classmethod, staticmethod)
411            ):
412                setattr(clsFnMethod, key, Enforcer(value, enabled=enabled))
413        return clsFnMethod
414    else:
415        raise Exception(
416            "Enforcer can only be used on classes, methods, or functions."
417        )
418
419
420Enforcer = Partial(Enforcer, enabled=True)