Get type object defined inside doctest by type name

I am trying to doc-test a method that accepts a module object module and a string with the name of the type type_name inside that module :

def get_type_from_string(module, type_name):
    """
    >>> class A: pass
    >>> A.__module__
    'mymodule'
    >>> get_type_from_string(sys.modules['mymodule'], 'A')
    <class '__mymodule__.A'>    <------ fails
    """
    return getattr(module, type_name)

When I am trying to get the type object with getattr (in reality, the method does more than just that), I am getting the error:

AttributeError: module 'mymodule' has no attribute 'A' 

Is it possible to doc-test this method without having to define the A class outside of the doc-test?

Answer

For some reason the defined A class is not present in the __dict__. Maybe this is a problem with the doctest itself. As a work around you can add it like this:

sys.modules[A.__module__].A = A

or

sys.modules[A.__module__].__dict__['A'] = A

def get_type_from_string(module, type_name):
    """
    >>> import sys
    >>> class A: pass
    >>> A.__module__
    'mymodule'
    >>> sys.modules[A.__module__].A = A
    >>> get_type_from_string(sys.modules['mymodule'], 'A')
    <class 'mymodule.A'>
    """
    return getattr(module, type_name)

This is how I am calling the test:

if __name__ == "__main__":
    # import sys
    import doctest
    import mymodule
    doctest.testmod(mymodule)