如何让Astroid通过类型注解推断函数参数的类型？

问题背景 #

我正尝试使用Astroid为Pylint编写自定义lint规则，但在类型推断方面遇到了问题。代码示例：

class ProtectedObject:
    def save():
        pass
                
def save_protected_object(obj: ProtectedObject) -> None:
    obj.save()

通过astroid.extract_node(code)成功获取到save_protected_object对应的FuncDef节点，但Astroid无法识别参数obj的类型注解，调用func.args.args[0].infer()返回Uninferable。

可行解决方案 #

方案1：启用全局类型注解支持 #

Astroid默认未开启类型注解的自动推断，需手动开启use_type_annotations配置：

import astroid

# 开启类型注解推断
astroid.MANAGER.config.use_type_annotations = True

code = """
class ProtectedObject:
    def save():
        pass
                
def save_protected_object(obj: ProtectedObject) -> None:
    obj.save()
"""

# 提取函数节点
func = astroid.extract_node(code + "\nsave_protected_object")
param = func.args.args[0]

# 推断参数类型
inferred_types = list(param.infer())
# 此时inferred_types会包含ProtectedObject对应的类节点

方案2：手动解析参数的注解节点 #

若全局配置不合适，可直接从参数节点的annotation属性获取注解，再单独推断注解的类型：

import astroid

code = """
class ProtectedObject:
    def save():
        pass
                
def save_protected_object(obj: ProtectedObject) -> None:
    obj.save()
"""

func = astroid.extract_node(code + "\nsave_protected_object")
param = func.args.args[0]

# 获取并推断注解类型
if param.annotation:
    inferred_annotation = list(param.annotation.infer())
    # inferred_annotation将返回ProtectedObject的类节点

对于泛型、Union这类复杂类型注解，需要额外处理逻辑，但简单类类型注解用上述方法即可解决问题。

内容的提问来源于stack exchange，提问作者unpollito