Skip to content

Commit

Permalink
Added AnnotatedTarget
Browse files Browse the repository at this point in the history
  • Loading branch information
evtn committed Dec 19, 2022
1 parent 49834bb commit 522da1e
Show file tree
Hide file tree
Showing 5 changed files with 52 additions and 7 deletions.
22 changes: 20 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -342,7 +342,7 @@ from gekkota import Assignment, Name

a, b, c = map(Name, "abc")

# Assignment(targets: Sequence[AssignmentTarget], value: Expression)
# Assignment(targets: Sequence[AssignmentTarget] | AnnotatedTarget, value: Expression)

print(
Assignment([a], b), # a = b
Expand All @@ -352,8 +352,26 @@ print(

```

For augmented assignment (e.g. `+=`) use `AugmentedAssignment`:
To annotate assignment (or just annotate a variable), use `AnnotatedTarget`:

```python

from gekkota import Assignment, AnnotatedTarget, Name

a, b, c = map(Name, "abc")
D = Name("D")

# AnnotatedTarget(target: AssignmentTarget, annotation: Expression)
print(
Assignment(AnnotatedTarget(a, D), b), # a: D = b
Assignment(AnnotatedTarget(a.index(b), D), c) # a[b]: D = c
Assignment([a, b], c), # a = b = c
)

```


For augmented assignment (e.g. `+=`) use `AugmentedAssignment`:

```python
from gekkota import Assignment, Name
Expand Down
1 change: 1 addition & 0 deletions gekkota/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
from .assignment import (
Assignment as Assignment,
AugmentedAssignment as AugmentedAssignment,
AnnotatedTarget as AnnotatedTarget,
)

from .control_flow import (
Expand Down
26 changes: 23 additions & 3 deletions gekkota/assignment.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from .utils import Utils
from .constants import Config, StrGen
from .core import Statement
from .core import Renderable, Statement
from .values import GetAttr, Indexing, Name


Expand All @@ -22,12 +22,32 @@


class Assignment(Statement):
def __init__(self, targets: Sequence[AssignmentTarget], value: Expression):
def __init__(
self, targets: Sequence[AssignmentTarget] | AnnotatedTarget, value: Expression
):
self.targets = targets
self.value = value

def render(self, config: Config) -> StrGen:
yield from Utils.separated(" = ", [*self.targets, self.value], config)
if isinstance(self.targets, AnnotatedTarget):
yield from self.targets.render(config)
else:
yield from Utils.separated(" = ", self.targets, config)
yield " "
yield "="
yield " "
yield from self.value.render(config)


class AnnotatedTarget(Statement):
def __init__(self, target: AssignmentTarget, annotation: Expression):
self.target = target
self.annotation = annotation

def render(self, config: Config) -> StrGen:
yield from self.target.render(config)
yield ": "
yield from self.annotation.render(config)


class AugmentedAssignment(Statement):
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "gekkota"
version = "0.4.2"
version = "0.5.0"
description = "Python code-generation for Python"
authors = ["Dmitry Gritsenko <k01419q45@ya.ru>"]
license = "MIT"
Expand Down
8 changes: 7 additions & 1 deletion test/test_assignment.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from gekkota import Name
from gekkota import Assignment, AugmentedAssignment
from gekkota import Assignment, AugmentedAssignment, AnnotatedTarget


a = Name("a")
Expand Down Expand Up @@ -32,3 +32,9 @@ def test_assignment(self):
def test_augassign(self):
for op in aug_ops:
assert str(AugmentedAssignment(a, op, b)) == f"a {op} b"

def test_annasign(self):
ab = AnnotatedTarget(a, b)
assert str(ab) == "a: b"

assert str(Assignment(ab, c)) == "a: b = c"

0 comments on commit 522da1e

Please sign in to comment.