CODEWITHSUNDEEP
pythonpython-3.xintrospection
Nick Gotch
Nick Gotch

Reputation: 9407

Get all the class types referenced within a function or class

I'm working on a code analyzer and I'm trying to identify all the class types referenced within a function or class in Python.

For example say I have this class:

import collections

Bar = collections.namedtuple('Bar', ['bar'])
Baz = collections.namedtuple('Baz', ['baz'])

class Foo(object):
    def get_bar(self):
        return Bar(1)
    def get_baz(self):
        return Baz(2)

I'm looking for a way I can get types of the functions and classes. Something like this:

print(get_types(Foo.get_bar)) # ['Bar']
print(get_types(Foo.get_baz)) # ['Baz']
print(get_types(Foo)) # ['Bar','Baz']

Upvotes: 1

Views: 95

Answers (1)

houtanb
houtanb

Reputation: 4100

One solution could involve using type annotations. Setting the return value of get_bar() to Bar and the return value of get_baz() to Baz, you could write get_types() as below...

import inspect
import collections

Bar = collections.namedtuple('Bar', ['bar'])
Baz = collections.namedtuple('Baz', ['baz'])


class Foo(object):
    def get_bar(self) -> Bar:
        return Bar(1)
    def get_baz(self) -> Baz:
        return Baz(2)


def get_types(obj):
    if inspect.isclass(obj):
        methods = [method for method in dir(obj)
                   if callable(getattr(obj, method)) and not method.startswith('_')]
        return [get_types(getattr(obj, method)) for method in methods]

    if callable(obj):
        return [obj.__annotations__['return'].__name__]


def main():
    print(get_types(Foo.get_bar)) # ['Bar']
    print(get_types(Foo.get_baz)) # ['Baz']
    print(get_types(Foo)) # ['Bar','Baz']


if __name__ == '__main__':
    main()

In get_types(obj), if obj contains an instance of a class, you can select out the non private methods of that class and return get_types() on each of these. If obj contains a function, then we just return the return attribute of that function.

Upvotes: 1

Related Questions