Add set comprehension support (#3)

* Add SetComp support and tests

* Update src/converter/code_generator_fixed.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: CrazyDubya

src/analyzer/code_analyzer_fixed.py

@@ -485,6 +485,10 @@ def _infer_expression_type(self, node: ast.AST) -> str:
                 elt_type = self._infer_expression_type(node.elts[0])
                 return f'std::set<{elt_type}>'
             return 'std::set<int>'
+        elif isinstance(node, ast.SetComp):
+            # Infer type from the element expression of the comprehension
+            elt_type = self._infer_expression_type(node.elt)
+            return f'std::set<{elt_type}>'
         elif isinstance(node, ast.Tuple):
             if node.elts:
                 elt_types = []

src/converter/code_generator_fixed.py

@@ -1104,6 +1104,28 @@ def _translate_expression(self, node: ast.AST, local_vars: Dict[str, str]) -> st
                 value_type = self._infer_cpp_type(node.values[0], local_vars)
 
             return f"std::map<{key_type}, {value_type}>{{{', '.join(pairs)}}}"
+        elif isinstance(node, ast.SetComp):
+            # Translate set comprehension using a lambda that fills a std::set
+            comp = node.generators[0]
+            iter_expr = self._translate_expression(comp.iter, local_vars)
+            target = self._translate_expression(comp.target, local_vars)
+            element_expr = self._translate_expression(node.elt, local_vars)
+            elem_type = self._infer_cpp_type(node.elt, local_vars)
+            conditions = ''
+            if comp.ifs:
+                conds = ' && '.join(f"({self._translate_expression(c, local_vars)})" for c in comp.ifs)
+                conditions = f"if ({conds}) "
+
+            lines = [
+                "[&]() {",
+                f"    std::set<{elem_type}> _set;",
+                f"    for (auto {target} : {iter_expr}) {{",
+                f"        {conditions} _set.insert({element_expr});",
+                "    }",
+                "    return _set;",
+                "}()",
+            ]
+            return "\n".join(lines)
         elif isinstance(node, ast.Tuple):
             # Handle tuple literals
             elements = [self._translate_expression(elt, local_vars) for elt in node.elts]
@@ -1257,6 +1279,15 @@ def _infer_cpp_type(self, node: ast.AST, local_vars: Dict[str, str]) -> str:
                 return f"std::map<{key_type}, {value_type}>"
             else:
                 return "std::map<std::string, int>"
+        elif isinstance(node, ast.Set):
+            if node.elts:
+                element_type = self._infer_cpp_type(node.elts[0], local_vars)
+                return f"std::set<{element_type}>"
+            else:
+                return "std::set<int>"
+        elif isinstance(node, ast.SetComp):
+            element_type = self._infer_cpp_type(node.elt, local_vars)
+            return f"std::set<{element_type}>"
         elif isinstance(node, ast.Tuple):
             if node.elts:
                 element_types = [self._infer_cpp_type(elt, local_vars) for elt in node.elts]

tests/test_code_analyzer_fixed.py

@@ -310,6 +310,24 @@ def test_inference_expressions(self):
             values=[ast.Constant(value=True), ast.Constant(value=False)]
         )
         assert analyzer._infer_expression_type(bool_op) == 'bool'
+
+    def test_set_comprehension_inference(self):
+        """Ensure set comprehensions are inferred as std::set."""
+        analyzer = CodeAnalyzer()
+
+        comp = ast.SetComp(
+            elt=ast.Name(id='x', ctx=ast.Load()),
+            generators=[
+                ast.comprehension(
+                    target=ast.Name(id='x', ctx=ast.Store()),
+                    iter=ast.Call(func=ast.Name(id='range', ctx=ast.Load()), args=[ast.Constant(value=5)], keywords=[]),
+                    ifs=[],
+                    is_async=0
+                )
+            ]
+        )
+
+        assert analyzer._infer_expression_type(comp) == 'std::set<int>'
 
     def test_type_annotation_handling(self):
         """Test handling of Python type annotations."""

tests/test_conversion_fixed.py

@@ -1,5 +1,7 @@
 import pytest
 from pathlib import Path
+import tempfile
+import os
 from src.analyzer.code_analyzer_fixed import CodeAnalyzer
 from src.rules.rule_manager import RuleManager
 from src.rules.basic_rules import (
@@ -56,4 +58,40 @@ def test_fibonacci_conversion(tmp_path):
     # Verify CMake content
     cmake_content = (output_dir / "CMakeLists.txt").read_text()
     assert "cmake_minimum_required" in cmake_content
-    assert "project(pytocpp_generated)" in cmake_content
+    assert "project(pytocpp_generated)" in cmake_content
+
+
+def test_set_comprehension_translation(tmp_path):
+    analyzer = CodeAnalyzer()
+    rule_manager = RuleManager()
+
+    rule_manager.register_rule(VariableDeclarationRule())
+    rule_manager.register_rule(FunctionDefinitionRule())
+    rule_manager.register_rule(ClassDefinitionRule())
+
+    generator = CodeGenerator(rule_manager)
+
+    with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as temp:
+        temp.write(
+            "def make_set(n):\n"
+            "    return {i * 2 for i in range(n)}\n"
+        )
+        temp_path = Path(temp.name)
+
+    try:
+        analysis_result = analyzer.analyze_file(temp_path)
+
+        rule_manager.set_context({
+            'type_info': analysis_result.type_info,
+            'performance_bottlenecks': analysis_result.performance_bottlenecks,
+            'memory_usage': analysis_result.memory_usage,
+            'hot_paths': analysis_result.hot_paths
+        })
+
+        output_dir = tmp_path / "generated_set"
+        generator.generate_code(analysis_result, output_dir)
+
+        impl_content = (output_dir / "generated.cpp").read_text()
+        assert "std::set<int> _set" in impl_content
+    finally:
+        os.unlink(temp_path)