How do I annotate a type which has an extra attribute with mypy?

user_562769 0
Viewed 384

I have an ast.UnaryOp object, but I manually added a parent attribute (see answer). How do I annotate this in a function?

I currently simply have:

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

But mypy complains (rightfully so) that ast.UnaryOp does not have the parent attribute.

How can I tell mypy that the node is not ast.UnaryOp but ast.UnaryOp + parent attribute?

My Try

I've created my own UnaryOp class which has a parent attribute. I can use this for type-casting:

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

The downside of it is that I need to type cast in a lot of places and that I introduced Any. I would prefer if I could just state somewhere that all ast.* types in that file do have a parent attribute

type-annotation mypy python
1 Answers

As a workaround you could use isinstance() with protocol decorated with the @runtime_checkable, that will check at runtime that all protocol members are defined and also ensure statical correctness.

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()
user_137826690
Related