Skip to content

Commit

Permalink
transformer: support unary literals
Browse files Browse the repository at this point in the history
  • Loading branch information
KotlinIsland committed Nov 14, 2024
1 parent f713c87 commit f369e4d
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 1 deletion.
12 changes: 11 additions & 1 deletion basedtyping/transformer.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,9 @@ def _basedtyping(self, attr: str) -> ast.Attribute:
)
return ast.fix_missing_locations(result)

def _literal(self, value: ast.Constant | ast.Name | ast.Attribute) -> ast.Subscript:
def _literal(
self, value: ast.Constant | ast.Name | ast.Attribute | ast.UnaryOp
) -> ast.Subscript:
return self.subscript(self._typing("Literal"), value)

def subscript(self, value: ast.expr, slice_: ast.expr) -> ast.Subscript:
Expand Down Expand Up @@ -178,6 +180,14 @@ def visit_Constant(self, node: ast.Constant) -> ast.AST:
return self._literal(node)
return node

@override
def visit_UnaryOp(self, node: ast.UnaryOp) -> ast.AST:
if not isinstance(node.operand, ast.Constant):
return node
if not isinstance(node.op, (ast.UAdd, ast.USub)):
return node
return self._literal(node)

@override
def visit_Tuple(self, node: ast.Tuple) -> ast.AST:
with self.implicit_tuple(value=False):
Expand Down
5 changes: 5 additions & 0 deletions tests/test_transformer.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,11 @@ def test_literal():
validate("1 | 2", Union[Literal[1], Literal[2]])


def test_negative():
validate("-1", Literal[-1])
validate("+1", Literal[1])


def test_literal_union():
validate("Union[1, 2]", Union[Literal[1], Literal[2]])

Expand Down

0 comments on commit f369e4d

Please sign in to comment.