Как аннотировать тип с дополнительным атрибутом с помощью mypy?

У меня есть объект ast.UnaryOp, но я вручную добавил атрибут parent (см. ответ). Как аннотировать это в функции?

У меня сейчас просто есть:

def _get_sim206(node: ast.UnaryOp):
    if isinstance(node.parent, ast.If):
        return False
    return True

Но mypy жалуется (правильно), что ast.UnaryOp не имеет атрибута parent.

Как я могу сказать mypy, что node не ast.UnaryOp, а ast.UnaryOp + parent attribute?

Моя попытка

Я создал свой собственный класс UnaryOp, у которого есть родительский атрибут. Я могу использовать это для приведения типов:

class UnaryOp(ast.UnaryOp):
    def __init__(self, orig: ast.UnaryOp) -> None:
        self.op = orig.op
        self.operand = orig.operand
        self.lineno = orig.lineno
        self.col_offset = orig.col_offset
        self.parent: ast.Expr = orig.parent  # type: ignore

Недостатком этого является то, что мне нужно вводить приведение во многих местах, и я ввел Any. Я бы предпочел, чтобы я мог просто указать где-нибудь, что все типы ast.* в этом файле имеют атрибут parent


person Martin Thoma    schedule 04.01.2021    source источник


Ответы (1)


В качестве обходного пути вы можете использовать isinstance() с протоколом, украшенным @runtime_checkable, который будет проверять во время выполнения, что все члены протокола определены, а также обеспечивают статическую корректность.

from typing import Protocol, runtime_checkable, cast
import ast


@runtime_checkable
class ParentProto(Protocol):
    parent: ast.AST
    

def foo(node: ast.UnaryOp):
    if isinstance(node, ParentProto):
        # use node.parent
        p = node.parent
        l = node.lineno
# assignment
g = ast.UnaryOp()
cast(ParentProto, g).parent = ast.If()
person alex_noname    schedule 04.01.2021
comment
О, классно! Я всегда забываю о protocol - person Martin Thoma; 04.01.2021
comment
Можно ли указать, что параметр должен быть ast.UnaryOp и ParentProto? (не Union, это логическое ИЛИ типов) - person Martin Thoma; 04.01.2021
comment
Насколько я знаю, это невозможно, у mypy проблемы с представлением пересечения типов, в отличие от union. - person alex_noname; 04.01.2021