CODEWITHSUNDEEP
pythonpython-typing
BAKE ZQ
BAKE ZQ

Reputation: 807

How to get type annotation within python function scope?

For example:

def test():
    a: int
    b: str
    print(__annotations__)
test()

This function call raises a NameError: name '__annotations__' is not defined error.

What I want is to get the type annotation within the function test, like the returned dict of annotations in the global scope or class scope.

It there are any ways can achieve this ?

If it's not possiblle, why this syntax exists?

Upvotes: 9

Views: 6186

Answers (3)

Sarctiann
Sarctiann

Reputation: 11

This version of the BLAKE approach also captures more complex types

import inspect
import re

def get_types_annotations(fn):
    source = inspect.getsource(fn)
    var_tps = {
        k: v for k, v in re.findall(r"(\w+) ?: ?([\w\"\[\], ']+[^ ]) ?=", source)
    }

    return var_tps

Upvotes: 0

BAKE ZQ
BAKE ZQ

Reputation: 807

Another simple solution I copied from others.

import re
def fn(q: int):
    a: int = 1
def get_types(fn):
    source = inspect.getsource(fn)
    var_tps = re.findall("  +([a-z0-9]+) *?: *([a-z0-9]+) *=", source)
    return var_tps

get_types(fn) # [('a', 'int')]

Upvotes: 0

Martijn Pieters
Martijn Pieters

Reputation: 1124458

Within a function, annotations for local variables are not retained, and so can't be accessed within the function. Only annotations for variables at the module and class-level result in an __annotations__ object being attached.

From the PEP 526 specification:

Annotating a local variable will cause the interpreter to treat it as a local, even if it was never assigned to. Annotations for local variables will not be evaluated[.]

[...]

In addition, at the module or class level, if the item being annotated is a simple name, then it and the annotation will be stored in the __annotations__ attribute of that module or class[.]

The __annotations__ global is only set if there are actual module-level annotations defined; the data model states it is optional:

Modules
[...] Predefined (writable) attributes: [...]; __annotations__ (optional) is a dictionary containing variable annotations collected during module body execution; [...].

When defined, you can access it from functions within the module, or via the globals() function.

If you try this within a function inside a class statement, then know that the class body namespace is not part of the scope of nested functions:

The scope of names defined in a class block is limited to the class block; it does not extend to the code blocks of methods – this includes comprehensions and generator expressions since they are implemented using a function scope.

You'd instead access the class namespace via a reference to the class. You can get such a reference by using the class global name, or inside bound methods, via type(self), inside class methods via the cls argument. Just use ClassObject.__annotations__ in that case.

If you must have access to annotations in the function local body, you'll need to parse the source code yourself. The Python AST does retain local annotations:

>>> import ast
>>> mod = ast.parse("def foo():\n    a: int = 0")
>>> print(ast.dump(mod.body[0], indent=4))
FunctionDef(
    name='foo',
    args=arguments(
        posonlyargs=[],
        args=[],
        kwonlyargs=[],
        kw_defaults=[],
        defaults=[]),
    body=[
        AnnAssign(
            target=Name(id='a', ctx=Store()),
            annotation=Name(id='int', ctx=Load()),
            value=Constant(value=0),
            simple=1)],
    decorator_list=[])

The above shows a text representation for the body of a function with a single annotation; the AnnAssign node tells us that a is annotated as int. You could collect such annotations with:

import inspect
import ast

class AnnotationsCollector(ast.NodeVisitor):
    """Collects AnnAssign nodes for 'simple' annotation assignments"""

    def __init__(self):
        self.annotations = {}

    def visit_AnnAssign(self, node):
        if node.simple:
            # 'simple' == a single name, not an attribute or subscription.
            # we can therefore count on `node.target.id` to exist. This is
            # the same criteria used for module and class-level variable
            # annotations.
            self.annotations[node.target.id] = node.annotation

def function_local_annotations(func):
    """Return a mapping of name to string annotations for function locals

    Python does not retain PEP 526 "variable: annotation" variable annotations
    within a function body, as local variables do not have a lifetime beyond
    the local namespace. This function extracts the mapping from functions that
    have source code available.
 
    """
    source = inspect.getsource(func)
    mod = ast.parse(source)
    assert mod.body and isinstance(mod.body[0], (ast.FunctionDef, ast.AsyncFunctionDef))
    collector = AnnotationsCollector()
    collector.visit(mod.body[0])
    return {
        name: ast.get_source_segment(source, node)
        for name, node in collector.annotations.items()
    }

The above walker finds all AnnAssignment annotations in the source code for a function object (and so requires that there is a source file available), then uses the AST source line and column information to extract the annotation source.

Given your test function, the above produces:

>>> function_local_annotations(test)
{'a': 'int', 'b': 'str'}

The type hints are not resolved, they are just strings, so you'll still have to use the typing.get_type_hints() function to turn these annotations into type objects.

Upvotes: 15

Related Questions