pythonpython-importpython-typing

How to avoid class Foo from being treated as __main__.Foo?


I have these two classes Foo and Bar which are separated in two different files (foo.py and bar.py) for organization purposes. They both implement a __add__ method. I want both __add__ methods to work whether other is an instance of class Foo or Bar.

However, if foo is an instance of class Foo and bar is an instance of class Bar, and I attempt to add bar + foo in the module foo, Boo's __add__ method will treat foo as __main__.Foo and will raise TypeError. The code in question can be summerized like so:

from typing import Union, Any

import bar as b

FooBar = Union['Foo', 'b.Bar']


class Foo:
    message: str

    def __init__(self, message: str) -> None:
        self.message = message

    def __add__(self, other: Any) -> 'Foo':
        if not isinstance(other, (Foo, b.Bar)):
            raise TypeError('can only add Foo and Bar')
        return Foo(' '.join((self.message, other.message)))


def test():
    foo = Foo('Jalim')
    bar = b.Bar('Rabei')
    foobar = foo + bar # works just fine
    print(foobar.message)
    barfoo = bar + foo # raises TypeError because it receives __main__.Foo
    print(barfoo.message)


if __name__ == '__main__':
    test()

from typing import Union, Any

import foo as f

FooBar = Union['f.Foo', 'Bar']


class Bar:
    message: str

    def __init__(self, message: str) -> None:
        self.message = message

    def __add__(self, other: Any) -> 'Bar':
        if not isinstance(other, (f.Foo, Bar)):
            print(f'{type(other) = }')
            raise TypeError('can only add Foo and Bar')
        return Bar(''.join((self.message, other.message)))

I know I could run test() in another file, but I'd like to be able to run from foo.py and while circunventing, if possible, the conversion from Foo to __main__.Foo.

At first, I tried using the type alias FooBar instead of a tuple (Foo, bar.Bar). However, since FooBar doens't really hold types and, instead, holds ForwardRefs, it can't be used in isinstance. I can't put Foo and Bar instead of 'Foo' and 'b.Bar' into a Union either, because it would generate a circular dependency issue.


Solution

  • With this structure, your class Foo is going to be created twice. It will exist with dual identities in two separate modules, and with two different memory locations:

    sys.modules["foo"].Foo
    sys.modules["__main__"].Foo
    

    For all practical purposes, they are different types, even though they have the same source code.

    The foo.Foo is created first, when the line import bar as b was executed, which itself triggers import foo as f. Then __main__.Foo is created, when the class definition is again encountered during the execution of the script in the top-level code environment.

    The solution is to separate the test code from the library code.

    Library code:

    # foo.py
    class Foo:
        ...
    
    # bar.py
    class Bar:
        ...
    

    Test code:

    from foo import Foo
    from bar import Bar
    
    def test_stuff():
        ...
    

    With this pattern, your top-level code environment will be a test script (or, more commonly, a test runner's entry-point) and the Foo type will only exist once, in the foo module namespace.