diff --git a/README.md b/README.md index 24c4373..832f57d 100644 --- a/README.md +++ b/README.md @@ -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 @@ -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 diff --git a/gekkota/__init__.py b/gekkota/__init__.py index 87e802d..432a768 100644 --- a/gekkota/__init__.py +++ b/gekkota/__init__.py @@ -46,6 +46,7 @@ from .assignment import ( Assignment as Assignment, AugmentedAssignment as AugmentedAssignment, + AnnotatedTarget as AnnotatedTarget, ) from .control_flow import ( diff --git a/gekkota/assignment.py b/gekkota/assignment.py index 82ebfc7..296f648 100644 --- a/gekkota/assignment.py +++ b/gekkota/assignment.py @@ -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 @@ -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): diff --git a/pyproject.toml b/pyproject.toml index d217da0..ad918cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 "] license = "MIT" diff --git a/test/test_assignment.py b/test/test_assignment.py index b71e8ba..043049f 100644 --- a/test/test_assignment.py +++ b/test/test_assignment.py @@ -1,5 +1,5 @@ from gekkota import Name -from gekkota import Assignment, AugmentedAssignment +from gekkota import Assignment, AugmentedAssignment, AnnotatedTarget a = Name("a") @@ -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"