Reference local type in pycontract

Question

I'm trying to use PyContracts within a web application, so I have lots of custom-defined classes being passed around that I simply want to type check alongside other more traditional argument types. I'd like to use contractual programming (PyContracts) to accomplish this, for the sake of cleanliness and forced documentation.

When I reference a locally visible class by name, PyContracts doesn't seem to be aware of the type. For example:

from contracts import contract

class SomeClass:
    pass

@contract
def f(a):
    """

    :param a: Just a parameter
    :type a: SomeClass
    """
    print(a)

my_a = SomeClass()
f(my_a)

Raises the following error:

ContractSyntaxError: Unknown identifier 'SomeClass'. Did you mean 'np_complex64'? (at char 0), (line:1, col:1)

I know I can use new_contract to custom-define names and bind them to classes, but that's a lot of hassle to do for every type. I want to use the docstring syntax for PyContracts if at all possible, and I definitely need to use the string-defined contract format since I'm using boolean type logic ("None|str|SomeClass"). How do I accomplish this with local types and minimal intrusion into the rest of my codebase?

Answer 1

I hacked together a magic decorator that adds types before creating the contract. For anyone that's interested, it seems to work, but it's probably pretty slow if you define a large number of functions:

def magic_contract(*args, **kwargs):
    # Check if we got called without arguments, just the function
    func = None
    if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
        func = args[0]
        args = tuple()

    def inner_decorator(f):
        for name, val in f.__globals__.items():
            if isinstance(val, type):
                new_contract(name, val)
        return contract(*args, **kwargs)(f)

    if func:
        return inner_decorator(func)
    return inner_decorator

And some test runs:

In [3]: class SomeClass:
   ...:     pass
   ...:

In [4]: @magic_contract
   ...: def f(a):
   ...:     """
   ...:
   ...:     :param a: Some parameter
   ...:     :type a: None|SomeClass
   ...:     """
   ...:     print(a)
   ...:

In [5]: f(None)
None

In [6]: f(SomeClass())
<__main__.SomeClass object at 0x7f1fa17c8b70>

In [7]: f(2)
...
ContractNotRespected: Breach for argument 'a' to f().
...

In [8]: @magic_contract(b='int|SomeClass')
   ...: def g(b):
   ...:     print(type(b))
   ...:

In [9]: g(2)
<class 'int'>

In [10]: g(SomeClass())
<class '__main__.SomeClass'>

In [11]: g(None)
...
ContractNotRespected: Breach for argument 'b' to g().
...