Added functionality for set notation (#8)

Author: wx1424

lf_toolkit/parse/set/__init__.py

@@ -15,3 +15,4 @@
 from .printer import UnicodePrinter
 from .transformer import SymPyBooleanTransformer
 from .transformer import SymPyTransformer
+from .ast import SetNotation

lf_toolkit/parse/set/ast.py

@@ -1,6 +1,7 @@
 from abc import ABC
 from abc import abstractmethod
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from typing import Tuple
 
 
 @dataclass
@@ -58,6 +59,9 @@ class SymmetricDifference(BinaryOp):
 class Term(Set):
     value: str
 
+@dataclass
+class SetNotation(Set):
+    elements: Tuple[str, ...] = field(default_factory=tuple)
 
 class Universe(Set):
     pass
@@ -80,6 +84,9 @@ def transform(self, node: Set):
             transformed_children.append(self.transform(node.left))
             transformed_children.append(self.transform(node.right))
 
+        if isinstance(node, SetNotation):
+            transformed_children.append(node.elements)
+
         # Dispatch to the specific transformation method based on node type
         method_name = type(node).__name__
         transformer = getattr(self, method_name, self.__unhandled_node)

lf_toolkit/parse/set/grammar/ascii.lark

@@ -22,7 +22,10 @@ start: expression
 
 ?group: "(" expression ")" -> group
 
-?term: ID | universe
+?term: ID | universe | set_notation
+
+set_notation: "{" (ELEMENT ("," ELEMENT)*)? "}"
 
 ID: /[A-Z]/
 universe: "Ω" | "Omega"
+ELEMENT: /[A-Za-z0-9]+/

lf_toolkit/parse/set/parser.py

@@ -17,6 +17,7 @@
 from .ast import Term
 from .ast import Union
 from .ast import Universe
+from .ast import SetNotation
 
 
 class ParseError(Exception):
@@ -105,3 +106,6 @@ def universe(self, _):
 
     def group(self, items):
         return Group(items[1] if self.latex else items[0])
+
+    def set_notation(self, items):
+        return SetNotation(tuple(str(i) for i in items))

lf_toolkit/parse/set/printer.py

@@ -1,5 +1,5 @@
 from .ast import Set
-from .ast import SetTransformer
+from .ast import SetTransformer, SetNotation
 
 
 class LatexPrinter(SetTransformer):
@@ -30,6 +30,9 @@ def Term(self, value):
     def Universe(self):
         return "\\Omega"
 
+    def SetNotation(self, elements):
+        return "\\{" + ",".join(elements) + "\\}"
+
 
 class ASCIIPrinter(SetTransformer):
     def print(self, node: Set):
@@ -59,6 +62,9 @@ def Term(self, value):
     def Universe(self):
         return "Omega"
 
+    def SetNotation(self, elements):
+        return "{" + ",".join(elements) + "}"
+
 
 class UnicodePrinter(SetTransformer):
     def print(self, node: Set):

lf_toolkit/parse/set/transformer.py

@@ -9,8 +9,9 @@
 from sympy import Union
 from sympy import UniversalSet
 from sympy import Xor
+from sympy import Integer
 
-from .ast import SetTransformer
+from .ast import SetTransformer, SetNotation
 
 
 class SymPyTransformer(SetTransformer):
@@ -43,6 +44,9 @@ def Term(self, expr):
     def Universe(self):
         return UniversalSet
 
+    def SetNotation(self, elements):
+        return FiniteSet(*[Integer(e) if str(e).isdigit() else Symbol(e) for e in elements])
+
 
 class SymPyBooleanTransformer(SetTransformer):
     def __init__(self):