extended_ast.py

import ast
import typing as t


class ExDict(ast.Dict):

    def __init__(
            self,
            keys: t.List[ast.expr],
            values: t.List[ast.expr],
            ctx: t.Union[ast.Store, ast.Load],
            lineno=None,
            col_offset=None,
    ):
        super().__init__()
        self.keys = keys
        self.values = values
        self.ctx = ctx
        if lineno:
            self.lineno = lineno
        if col_offset:
            self.col_offset = col_offset

    _fields = ('keys', 'values', 'ctx')


class ConstantMapping(ast.Constant):

    def __init__(self, value, lineno=None, col_offset=None):
        super().__init__()
        self.value = value
        self.lineno = lineno
        self.col_offset = col_offset


class AssignExpr(ast.Assign):

    def __init__(self,
                 target: ast.expr,
                 value: ast.expr,
                 lineno: int = None,
                 col_offset: int = None):
        super().__init__()
        self.value = value
        self.targets = [target]
        self.lineno = lineno
        self.col_offset = col_offset

    @property
    def target(self):
        return self.targets[0]

    @target.setter
    def target(self, value):
        self.targets[0] = value