Add support for Collection type

src/astx/literals/__init__.py

@@ -34,6 +34,10 @@
     LiteralTime,
     LiteralTimestamp,
 )
+from astx.literals.collection import (
+    LiteralSet,
+    LiteralList,
+)
 
 __all__ = [
     "Literal",
@@ -60,4 +64,6 @@
     "LiteralTimestamp",
     "LiteralUTF8Char",
     "LiteralUTF8String",
+    "LiteralSet",
+    "LiteralList",
 ]

src/astx/literals/collection.py

@@ -0,0 +1,72 @@
+"""ASTx Data Types module."""
+
+from __future__ import annotations
+
+from public import public
+from typeguard import typechecked
+from typing import Any
+
+from astx.base import (
+    NO_SOURCE_LOCATION,
+    ReprStruct,
+    SourceLocation,
+)
+from astx.literals.base import Literal
+from astx.types.collection import (
+    Set,
+    List,
+)
+
+
+@public
+@typechecked
+class LiteralSet(Literal):
+    """LiteralSet data type class."""
+
+    value: Set[Any]
+
+    def __init__(
+        self, value: Set[Any], loc: SourceLocation = NO_SOURCE_LOCATION
+    ) -> None:
+        """Initialize LiteralSet."""
+        super().__init__(loc)
+        self.value = value
+        self.type_ = Set
+        self.loc = loc
+
+    def __str__(self) -> str:
+        """Return a string representation of the LiteralSet."""
+        return f"LiteralSet({self.value})"
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST representation for the Set literal."""
+        key = "Set"
+        value = list(self.value)
+        return self._prepare_struct(key, value, simplified)
+
+
+@public
+@typechecked
+class LiteralList(Literal):
+    """LiteralList data type class."""
+
+    value: List[Any]
+
+    def __init__(
+        self, value: List[Any], loc: SourceLocation = NO_SOURCE_LOCATION
+    ) -> None:
+        """Initialize LiteralList."""
+        super().__init__(loc)
+        self.value = value
+        self.type_ = List
+        self.loc = loc
+
+    def __str__(self) -> str:
+        """Return a string representation of the LiteralList."""
+        return f"LiteralList({self.value})"
+
+    def get_struct(self, simplified: bool = False) -> ReprStruct:
+        """Return the AST representation for the List literal."""
+        key = "List"
+        value = self.value
+        return self._prepare_struct(key, value, simplified)

src/astx/types/__init__.py

@@ -46,6 +46,10 @@
     Time,
     Timestamp,
 )
+from astx.types.collection import (
+    Set,
+    List,
+)
 
 __all__ = [
     "AnyType",
@@ -80,4 +84,6 @@
     "UTF8String",
     "UnsignedInteger",
     "TypeCastExpr",
+    "Set",
+    "List",
 ]

src/astx/types/collection.py

@@ -0,0 +1,26 @@
+"""ASTx Data Types module."""
+
+from __future__ import annotations
+
+from public import public
+from typeguard import typechecked
+
+from astx.types.base import AnyType
+
+
+@public
+@typechecked
+class Collection(AnyType):
+    """Base class for collection data types."""
+
+
+@public
+@typechecked
+class Set(Collection):
+    """Set data type expression."""
+
+
+@public
+@typechecked
+class List(Collection):
+    """List data type expression."""

tests/datatypes/test_collection.py

@@ -0,0 +1,133 @@
+"""Tests for List and Set data types."""
+
+from __future__ import annotations
+
+from typing import Callable, Type
+
+import astx
+import pytest
+
+from astx.types.operators import BinaryOp
+from astx.variables import Variable
+
+VAR_A = Variable("a")
+
+LIST_LITERAL_CLASSES = [
+    astx.LiteralList,
+]
+
+SET_LITERAL_CLASSES = [
+    astx.LiteralSet,
+]
+
+
+def test_variable() -> None:
+    """Test variable for List and Set."""
+    var_a = Variable("a")
+    var_b = Variable("b")
+
+    BinaryOp(op_code="+", lhs=var_a, rhs=var_b)
+
+
+@pytest.mark.parametrize("literal_class", LIST_LITERAL_CLASSES)
+def test_list_literal(literal_class: Type[astx.Literal]) -> None:
+    """Test List literals."""
+    lit_a = literal_class([1, 2, 3])
+    lit_b = literal_class([4, 5, 6])
+    BinaryOp(op_code="+", lhs=lit_a, rhs=lit_b)
+
+
+@pytest.mark.parametrize("literal_class", SET_LITERAL_CLASSES)
+def test_set_literal(literal_class: Type[astx.Literal]) -> None:
+    """Test Set literals."""
+    lit_a = literal_class({1, 2, 3})
+    lit_b = literal_class({4, 5, 6})
+    BinaryOp(op_code="+", lhs=lit_a, rhs=lit_b)
+
+
+@pytest.mark.parametrize(
+    "fn_bin_op,op_code",
+    [
+        (lambda literal_class: VAR_A + literal_class([1, 2, 3]), "+"),
+        (lambda literal_class: VAR_A == literal_class([1, 2, 3]), "=="),
+        (lambda literal_class: VAR_A != literal_class([1, 2, 3]), "!="),
+    ],
+)
+@pytest.mark.parametrize("literal_class", LIST_LITERAL_CLASSES)
+def test_bin_ops_list(
+    literal_class: Type[astx.Literal],
+    fn_bin_op: Callable[[Type[astx.Literal]], BinaryOp],
+    op_code: str,
+) -> None:
+    """Test binary operations on List literals."""
+    bin_op = fn_bin_op(literal_class)
+    assert bin_op.op_code == op_code
+    assert str(bin_op) != ""
+    assert repr(bin_op) != ""
+    assert bin_op.get_struct() != {}
+    assert bin_op.get_struct(simplified=True) != {}
+
+
+@pytest.mark.parametrize(
+    "fn_bin_op,op_code",
+    [
+        (lambda literal_class: VAR_A + literal_class({1, 2, 3}), "+"),
+        (lambda literal_class: VAR_A == literal_class({1, 2, 3}), "=="),
+        (lambda literal_class: VAR_A != literal_class({1, 2, 3}), "!="),
+    ],
+)
+@pytest.mark.parametrize("literal_class", SET_LITERAL_CLASSES)
+def test_bin_ops_set(
+    literal_class: Type[astx.Literal],
+    fn_bin_op: Callable[[Type[astx.Literal]], BinaryOp],
+    op_code: str,
+) -> None:
+    """Test binary operations on Set literals."""
+    bin_op = fn_bin_op(literal_class)
+    assert bin_op.op_code == op_code
+    assert str(bin_op) != ""
+    assert repr(bin_op) != ""
+    assert bin_op.get_struct() != {}
+    assert bin_op.get_struct(simplified=True) != {}
+
+
+@pytest.mark.parametrize(
+    "fn_unary_op,op_code",
+    [
+        (lambda literal_class: +literal_class([1, 2, 3]), "+"),
+    ],
+)
+@pytest.mark.parametrize("literal_class", LIST_LITERAL_CLASSES)
+def test_unary_ops_list(
+    literal_class: Type[astx.Literal],
+    fn_unary_op: Callable[[Type[astx.Literal]], BinaryOp],
+    op_code: str,
+) -> None:
+    """Test unary operations on List literals."""
+    unary_op = fn_unary_op(literal_class)
+    assert unary_op.op_code == op_code
+    assert str(unary_op) != ""
+    assert repr(unary_op) != ""
+    assert unary_op.get_struct() != {}
+    assert unary_op.get_struct(simplified=True) != {}
+
+
+@pytest.mark.parametrize(
+    "fn_unary_op,op_code",
+    [
+        (lambda literal_class: +literal_class({1, 2, 3}), "+"),
+    ],
+)
+@pytest.mark.parametrize("literal_class", SET_LITERAL_CLASSES)
+def test_unary_ops_set(
+    literal_class: Type[astx.Literal],
+    fn_unary_op: Callable[[Type[astx.Literal]], BinaryOp],
+    op_code: str,
+) -> None:
+    """Test unary operations on Set literals."""
+    unary_op = fn_unary_op(literal_class)
+    assert unary_op.op_code == op_code
+    assert str(unary_op) != ""
+    assert repr(unary_op) != ""
+    assert unary_op.get_struct() != {}
+    assert unary_op.get_struct(simplified=True) != {}