Python patching __new__ method

Viewed 454

I am trying to patch __new__ method of a class, and it is not working as I expect.

from contextlib import contextmanager

class A:
    def __init__(self, arg):
        print('A init', arg)

@contextmanager
def patch_a():
    new = A.__new__

    def fake_new(cls, *args, **kwargs):
        print('call fake_new')
        return new(cls, *args, **kwargs) 
        # here I get error: TypeError: object.__new__() takes exactly one argument (the type to instantiate)

    A.__new__ = fake_new
    try:
        yield
    finally:
        A.__new__ = new

if __name__ == '__main__':
    A('foo')
    with patch_a():
        A('bar')
    A('baz')

I expect the following output:

A init foo
call fake_new
A init bar
A init baz

But after call fake_new I get an error (see comment in the code). For me It seems like I just decorate a __new__ method and propagate all args unchanged. It doesn't work and the reason is obscure for me.

Also I can write return new(cls) and call A('bar') works fine. But then A('baz') breaks.

Can someone explain what is going on?

Python version is 3.8

1 Answers

You've run into a complicated part of Python object instantiation - in which the language opted for a design that would allow one to create a custom __init__ method with parameters, without having to touch __new__.

However, the in the base of class hierarchy, object, both __new__ and __init__ take one single parameter each.

IIRC, it goes this way: if your class have a custom __init__ and you did not touch __new__ and there are more any parameters to the class instantiation that would be passed to both __init__ and __new__, the parameters will be stripped from the call do __new__, so you don't have to customize it just to swallow the parameters you consume in __init__. The converse is also true: if your class have a custom __new__ with extra parameters, and no custom __init__, these are not passed to object.__init__.

With your design, Python sees a custom __new__ and passes it the same extra arguments that are passed to __init__ - and by using *args, **kw, you forward those to object.__new__ which accepts a single parameter - and you get the error you presented us.

The fix is to not pass those extra parameters to the original __new__ method - unless they are needed there - so you have to make the same check Python's type does when initiating an object.

And an interesting surprise to top it: while making the example work, I found out that even if A.__new__ is deleted when restoring the patch, it is still considered as "touched" by cPython's type instantiation, and the arguments are passed through. In order to get your code working I needed to leave a permanent stub A.__new__ that will forward only the cls argument:


from contextlib import contextmanager

class A:
    def __init__(self, arg):
        print('A init', arg)

@contextmanager
def patch_a():
    new = A.__new__

    def fake_new(cls, *args, **kwargs):
        print('call fake_new')
        if new is object.__new__:
            return new(cls)
        return new(cls, *args, **kwargs)
        # here I get error: TypeError: object.__new__() takes exactly one argument (the type to instantiate)

    A.__new__ = fake_new
    try:
        yield
    finally:
        del A.__new__
        if new is not object.__new__:
            A.__new__ = new
        else:
            A.__new__ = lambda cls, *args, **kw: object.__new__(cls)

        print(A.__new__)

if __name__ == '__main__':
    A('foo')
    with patch_a():
        A('bar')
    A('baz')

(I tried inspecting the original __new__ signature instead of the new is object.__new__ comparison - to no avail: object.__new__ signature is *args, **kwargs - possibly made so that it will never fail on static checking)

Related