feat: Add augmented assignment feature (#208)

Author: ahmed44374

docs/contributing.md

@@ -4,7 +4,7 @@ In order to be able to contribute, it is important that you understand the
 project layout. This project uses the _src layout_, which means that the package
 code is located at `./src/astx`.
 
-For my information, check the official documentation:
+For more information, check the official documentation:
 https://packaging.python.org/en/latest/discussions/src-layout-vs-flat-layout/
 
 In addition, you should know that to build our package we use

src/astx/__init__.py

@@ -115,6 +115,7 @@
 )
 from astx.operators import (
     AssignmentExpr,
+    AugAssign,
     VariableAssignment,
     WalrusOp,
 )
@@ -204,6 +205,7 @@ def get_version() -> str:
     "AssignmentExpr",
     "AsyncForRangeLoopExpr",
     "AsyncForRangeLoopStmt",
+    "AugAssign",
     "AwaitExpr",
     "BinaryOp",
     "Block",

src/astx/base.py

@@ -103,6 +103,7 @@ class ASTKind(Enum):
     BinaryOpKind = -301
     WalrusOpKind = -302
     AssignmentExprKind = -303
+    AugmentedAssignKind = -304
 
     # functions
     PrototypeKind = -400

src/astx/operators.py

@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
-from typing import Iterable, Optional, cast
+from typing import Iterable, Literal, Optional, cast
 
 from public import public
+from typing_extensions import TypeAlias
 
 from astx.base import (
     NO_SOURCE_LOCATION,
@@ -13,6 +14,7 @@
     DataType,
     DictDataTypesStruct,
     Expr,
+    Identifier,
     ReprStruct,
     SourceLocation,
     StatementType,
@@ -128,3 +130,55 @@ def get_struct(self, simplified: bool = False) -> ReprStruct:
         key = str(self)
         value = self.value.get_struct(simplified)
         return self._prepare_struct(key, value, simplified)
+
+
+OpCodeAugAssign: TypeAlias = Literal[
+    "+=",
+    "-=",
+    "*=",
+    "/=",
+    "//=",
+    "%=",
+    "**=",
+    "&=",
+    "|=",
+    "^=",
+    "<<=",
+    ">>=",
+]
+
+
+@public
+@typechecked
+class AugAssign(DataType):
+    """AST class for augmented assignment."""
+
+    target: Identifier
+    op_code: OpCodeAugAssign
+    value: DataType
+
+    def __init__(
+        self,
+        target: Identifier,
+        op_code: OpCodeAugAssign,
+        value: DataType,
+        loc: SourceLocation = NO_SOURCE_LOCATION,
+    ) -> None:
+        super().__init__(loc=loc)
+        self.target = target
+        self.op_code = op_code
+        self.value = value
+        self.kind = ASTKind.AugmentedAssignKind
+
+    def __str__(self) -> str:
+        """Return a string that represents the augmented assignment object."""
+        return f"AugAssign[{self.op_code}]"
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST structure of the object."""
+        key = str(self)
+        value: ReprStruct = {
+            "target": self.target.get_struct(simplified),
+            "value": self.value.get_struct(simplified),
+        }
+        return self._prepare_struct(key, value, simplified)

src/astx/tools/transpilers/python.py

@@ -5,6 +5,7 @@
 from plum import dispatch
 
 import astx
+import astx.operators
 
 from astx.tools.typing import typechecked
 
@@ -520,6 +521,13 @@ def visit(self, node: astx.WalrusOp) -> str:
         """Handle Walrus operator."""
         return f"({self.visit(node.lhs)} := {self.visit(node.rhs)})"
 
+    @dispatch  # type: ignore[no-redef]
+    def visit(self, node: astx.AugAssign) -> str:
+        """Handle Augmented assign operator."""
+        target = self.visit(node.target)
+        value = self.visit(node.value)
+        return f"{target} {node.op_code} {value}"
+
     @dispatch  # type: ignore[no-redef]
     def visit(self, node: astx.WhileExpr) -> str:
         """Handle WhileExpr nodes."""

tests/test_operators.py

@@ -1,9 +1,17 @@
 """Tests for operators."""
 
+from typing import cast
+
 import astx
+import pytest
 
 from astx.literals.numeric import LiteralInt32
-from astx.operators import AssignmentExpr, VariableAssignment
+from astx.operators import (
+    AssignmentExpr,
+    AugAssign,
+    OpCodeAugAssign,
+    VariableAssignment,
+)
 from astx.variables import Variable
 from astx.viz import visualize
 
@@ -116,3 +124,31 @@ def test_not_op() -> None:
     assert op.get_struct()
     assert op.get_struct(simplified=True)
     visualize(op.get_struct())
+
+
+@pytest.mark.parametrize(
+    "operator, value",
+    [
+        (cast(OpCodeAugAssign, "+="), 10),
+        (cast(OpCodeAugAssign, "-="), 5),
+        (cast(OpCodeAugAssign, "*="), 3),
+        (cast(OpCodeAugAssign, "/="), 2),
+        (cast(OpCodeAugAssign, "//="), 2),
+        (cast(OpCodeAugAssign, "%="), 4),
+        (cast(OpCodeAugAssign, "**="), 2),
+        (cast(OpCodeAugAssign, "&="), 6),
+        (cast(OpCodeAugAssign, "|="), 3),
+        (cast(OpCodeAugAssign, "^="), 1),
+        (cast(OpCodeAugAssign, "<<="), 1),
+        (cast(OpCodeAugAssign, ">>="), 2),
+    ],
+)
+def test_aug_assign_operations(operator: OpCodeAugAssign, value: int) -> None:
+    """Test all augmented assignment operators using parametrize."""
+    var_x = astx.Identifier(value="x")
+    literal_value = LiteralInt32(value)
+    aug_assign = AugAssign(var_x, operator, literal_value)
+
+    assert str(aug_assign) == f"AugAssign[{operator}]"
+    assert aug_assign.get_struct()
+    assert aug_assign.get_struct(simplified=True)