docstring before function definition

src/know_lang_bot/code_parser/parser.py

@@ -14,7 +14,6 @@ LOG = FancyLogger(__name__)
 class ChunkType(str, Enum):
     CLASS = "class"
     FUNCTION = "function"
-    MODULE = "module"
     OTHER = "other"
 
 class CodeChunk(BaseModel):
@@ -44,14 +43,36 @@ class CodeParser:
         self.language = Language(tree_sitter_python.language())
         self.parser = Parser(self.language)
 
-    def _extract_docstring(self, node: Node, source_code: bytes) -> Optional[str]:
-        """Extract docstring from a class or function node"""
-        for child in node.children:
-            if child.type == "expression_statement":
-                string_node = child.children[0]
-                if string_node.type in ("string", "string_literal"):
-                    return source_code[string_node.start_byte:string_node.end_byte].decode('utf-8')
-        return None
+    def _get_preceding_docstring(self, node: Node, source_code: bytes) -> Optional[str]:
+        """Extract docstring from comments"""
+        docstring_parts = []
+        current_node : Node = node.prev_sibling
+
+        while current_node:
+            print(current_node.text)
+            if current_node.type == "comment":
+                comment = source_code[current_node.start_byte:current_node.end_byte].decode('utf-8')
+                docstring_parts.insert(0, comment)
+            elif current_node.type == "expression_statement":
+                string_node = current_node.children[0] if current_node.children else None
+                if string_node and string_node.type in ("string", "string_literal"):
+                    docstring = source_code[string_node.start_byte:string_node.end_byte].decode('utf-8')
+                    docstring_parts.insert(0, docstring)
+                    
+                break
+            elif current_node.type not in ("empty_statement", "newline"):
+                break
+            current_node = current_node.prev_sibling
+        
+        return '\n'.join(docstring_parts) if docstring_parts else None
+
+    def _has_syntax_error(self, node: Node) -> bool:
+        """Check if the node or its children contain syntax errors"""
+        if node.type == "ERROR":
+            return True
+        if node.has_error:
+            return True
+        return any(self._has_syntax_error(child) for child in node.children)
 
     def parse_file(self, file_path: Path) -> List[CodeChunk]:
         """Parse a single file and return list of code chunks"""
@@ -64,6 +85,12 @@ class CodeParser:
                 source_code = f.read()
             
             tree = self.parser.parse(source_code)
+
+            # Check for overall syntax validity
+            if self._has_syntax_error(tree.root_node):
+                LOG.warning(f"Syntax errors found in {file_path}")
+                return []
+
             chunks: List[CodeChunk] = []
             
             # Process the syntax tree
@@ -73,15 +100,8 @@ class CodeParser:
                 elif node.type == "function_definition":
                     chunks.append(self._process_function(node, source_code, file_path))
                 else:
-                    # Store other top-level code as separate chunks
-                    if node.type not in ("comment", "empty_statement"):
-                        chunks.append(CodeChunk(
-                            type=ChunkType.OTHER,
-                            content=source_code[node.start_byte:node.end_byte].decode('utf-8'),
-                            start_line=node.start_point[0],
-                            end_line=node.end_point[0],
-                            file_path=str(file_path)
-                        ))
+                    # Skip other node types for now
+                    pass
             
             return chunks
         except Exception as e:
@@ -94,6 +114,9 @@ class CodeParser:
                    for child in node.children 
                    if child.type == "identifier")
         
+        if not name:
+            raise ValueError(f"Could not find class name in node: {node.text}")
+        
         return CodeChunk(
             type=ChunkType.CLASS,
             name=name,
@@ -101,7 +124,7 @@ class CodeParser:
             start_line=node.start_point[0],
             end_line=node.end_point[0],
             file_path=str(file_path),
-            docstring=self._extract_docstring(node, source_code)
+            docstring=self._get_preceding_docstring(node, source_code)
         )
 
     def _process_function(self, node: Node, source_code: bytes, file_path: Path) -> CodeChunk:
@@ -109,6 +132,9 @@ class CodeParser:
         name = next(child.text.decode('utf-8') 
                    for child in node.children 
                    if child.type == "identifier")
+
+        if not name:
+            raise ValueError(f"Could not find function name in node: {node.text}")
         
         return CodeChunk(
             type=ChunkType.FUNCTION,
@@ -117,7 +143,7 @@ class CodeParser:
             start_line=node.start_point[0],
             end_line=node.end_point[0],
             file_path=str(file_path),
-            docstring=self._extract_docstring(node, source_code)
+            docstring=self._get_preceding_docstring(node, source_code)
         )
 
     def parse_repository(self) -> List[CodeChunk]:

tests/test_constants.py

@@ -0,0 +1,114 @@
+from typing import Dict, NamedTuple
+
+class ExpectedChunk(NamedTuple):
+    name: str
+    docstring: str
+    content_snippet: str  # A unique part of the content that should be present
+
+# Test file contents
+SIMPLE_FUCNTION_CLEANED_DOCSTRING = "Say hello to the world"
+SIMPLE_FUNCTION_DOCSTRING = f'\"\"\"{SIMPLE_FUCNTION_CLEANED_DOCSTRING}\"\"\"'
+SIMPLE_FUNCTION = f'''
+{SIMPLE_FUNCTION_DOCSTRING}
+def hello_world():
+    return "Hello, World!"
+'''
+
+SIMPLE_CLASS_DOCSTRING = f'\"\"\"A simple class for testing\"\"\"'
+SIMPLE_CLASS = f'''
+{SIMPLE_CLASS_DOCSTRING}
+class SimpleClass:
+    def __init__(self):
+        self.value = 42
+
+    def get_value(self):
+        return self.value
+'''
+
+NESTED_OUTER_CLASS_DOCSTRING = "#Outer class docstring"
+NESTED_CLASS = f'''
+{NESTED_OUTER_CLASS_DOCSTRING}
+class OuterClass:
+    class InnerClass:
+        """Inner class docstring"""
+        def inner_method(self):
+            return "inner"
+
+    def outer_method(self):
+        return "outer"
+'''
+
+
+COMPLEX_FUNCTION_DOCSTRING = f"""\"\"\"
+A complex function with type hints and docstring
+Args:
+    param1: First parameter
+    param2: Optional second parameter
+Returns:
+    List of strings
+\"\"\""""
+
+COMPLEX_CLASS_DOCSTRING = "# # #Complex class implementation"
+COMPLEX_FILE = f'''
+import os
+from typing import List, Optional
+
+{COMPLEX_FUNCTION_DOCSTRING}
+def complex_function(param1: str, param2: Optional[int] = None) -> List[str]:
+    results = []
+    if param2 is not None:
+        results.extend([param1] * param2)
+    return results
+
+# Some comment
+CONSTANT = 42
+
+{COMPLEX_CLASS_DOCSTRING}
+class ComplexClass:
+    """Complex class implementation Test Test"""
+    def __init__(self):
+        self._value = None
+'''
+
+INVALID_SYNTAX = '''def invalid_syntax(:'''
+
+# Expected test results
+SIMPLE_FILE_EXPECTATIONS = {
+    'hello_world': ExpectedChunk(
+        name="hello_world",
+        docstring=SIMPLE_FUNCTION_DOCSTRING,
+        content_snippet='return "Hello, World!"'
+    ),
+    'SimpleClass': ExpectedChunk(
+        name="SimpleClass",
+        docstring=SIMPLE_CLASS_DOCSTRING,
+        content_snippet='self.value = 42'
+    )
+}
+
+NESTED_CLASS_EXPECTATIONS = {
+    'OuterClass': ExpectedChunk(
+        name="OuterClass",
+        docstring=NESTED_OUTER_CLASS_DOCSTRING,
+        content_snippet='class InnerClass'
+    )
+}
+
+COMPLEX_FILE_EXPECTATIONS = {
+    'complex_function': ExpectedChunk(
+        name="complex_function",
+        docstring=COMPLEX_FUNCTION_DOCSTRING,
+        content_snippet='List[str]'
+    ),
+    'ComplexClass': ExpectedChunk(
+        name="ComplexClass",
+        docstring=COMPLEX_CLASS_DOCSTRING,
+        content_snippet='_value = None'
+    )
+}
+
+TEST_FILES = {
+    'simple.py': SIMPLE_FUNCTION + SIMPLE_CLASS,
+    'nested.py': NESTED_CLASS,
+    'complex.py': COMPLEX_FILE
+}

tests/test_parser.py

@@ -0,0 +1,152 @@
+from know_lang_bot.code_parser.parser import CodeChunk, CodeParser, ChunkType
+from pathlib import Path
+from tests.test_constants import (
+    SIMPLE_FILE_EXPECTATIONS,
+    NESTED_CLASS_EXPECTATIONS,
+    COMPLEX_FILE_EXPECTATIONS,
+    INVALID_SYNTAX,
+    TEST_FILES,
+)
+import pytest
+import tempfile
+import git
+
+
+@pytest.fixture
+def temp_repo():
+    """Create a temporary git repository with sample Python files"""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Initialize git repo
+        repo = git.Repo.init(temp_dir)
+        
+        # Create sample Python files
+        for filename, content in TEST_FILES.items():
+            file_path = Path(temp_dir) / filename
+            file_path.write_text(content)
+            repo.index.add([str(file_path)])
+        
+        repo.index.commit("Initial commit")
+        
+        yield temp_dir
+
+def find_chunk_by_criteria(chunks: list[CodeChunk], **criteria) -> CodeChunk:
+    """Helper function to find a chunk matching given criteria"""
+    for chunk in chunks:
+        if all(getattr(chunk, k) == v for k, v in criteria.items()):
+            return chunk
+    return None
+
+def test_init_parser(temp_repo):
+    """Test parser initialization"""
+    parser = CodeParser(temp_repo)
+    assert parser.repo_path == Path(temp_repo)
+    assert parser.language is not None
+    assert parser.parser is not None
+
+def test_parse_simple_file(temp_repo):
+    """Test parsing a simple Python file with function and class"""
+    parser = CodeParser(temp_repo)
+    chunks = parser.parse_file(Path(temp_repo) / "simple.py")
+
+    # Test function
+    function_chunk = find_chunk_by_criteria(chunks, type=ChunkType.FUNCTION, name="hello_world")
+    assert function_chunk is not None
+    expected = SIMPLE_FILE_EXPECTATIONS['hello_world']
+    assert expected.content_snippet in function_chunk.content
+    assert function_chunk.docstring is not None
+    assert function_chunk.docstring in expected.docstring
+
+    # Test class
+    class_chunk = find_chunk_by_criteria(chunks, type=ChunkType.CLASS, name="SimpleClass")
+    assert class_chunk is not None
+    expected = SIMPLE_FILE_EXPECTATIONS['SimpleClass']
+    assert expected.content_snippet in class_chunk.content
+    assert class_chunk.docstring is not None
+    assert class_chunk.docstring in expected.docstring
+
+
+def test_parse_nested_classes(temp_repo):
+    """Test parsing nested class definitions"""
+    parser = CodeParser(temp_repo)
+    chunks = parser.parse_file(Path(temp_repo) / "nested.py")
+    
+    # Test outer class
+    outer_class = find_chunk_by_criteria(chunks, type=ChunkType.CLASS, name="OuterClass")
+    assert outer_class is not None
+    expected = NESTED_CLASS_EXPECTATIONS['OuterClass']
+    assert expected.content_snippet in outer_class.content
+    assert outer_class.docstring is not None
+    assert outer_class.docstring in expected.docstring
+
+    # Verify inner class: Not implemented yet
+    pass
+
+def test_parse_complex_file(temp_repo):
+    """Test parsing a complex Python file"""
+    parser = CodeParser(temp_repo)
+    chunks = parser.parse_file(Path(temp_repo) / "complex.py")
+    
+    # Test function with type hints
+    complex_func = find_chunk_by_criteria(
+        chunks, 
+        type=ChunkType.FUNCTION,
+        name="complex_function"
+    )
+    assert complex_func is not None
+    expected = COMPLEX_FILE_EXPECTATIONS['complex_function']
+    assert expected.content_snippet in complex_func.content
+    assert complex_func.docstring is not None
+    assert complex_func.docstring in expected.docstring
+    
+    # Test complex class
+    complex_class = find_chunk_by_criteria(
+        chunks,
+        type=ChunkType.CLASS,
+        name="ComplexClass"
+    )
+    assert complex_class is not None
+    expected = COMPLEX_FILE_EXPECTATIONS['ComplexClass']
+    assert expected.content_snippet in complex_class.content
+    assert complex_class.docstring is not None
+    assert complex_class.docstring in expected.docstring 
+
+
+def test_parse_repository(temp_repo):
+    """Test parsing entire repository"""
+    parser = CodeParser(temp_repo)
+    chunks = parser.parse_repository()
+    
+    file_paths = {chunk.file_path for chunk in chunks}
+    assert len(file_paths) == 3
+    
+    # Verify we can find chunks from each test file
+    for filename in TEST_FILES.keys():
+        file_chunks = [c for c in chunks if Path(c.file_path).name == filename]
+        assert len(file_chunks) > 0
+
+def test_error_handling(temp_repo):
+    """Test error handling for invalid files"""
+    parser = CodeParser(temp_repo)
+    
+    # Test invalid syntax
+    invalid_file = Path(temp_repo) / "invalid.py"
+    invalid_file.write_text(INVALID_SYNTAX)
+    chunks = parser.parse_file(invalid_file)
+    assert chunks == []
+    
+    # Test non-existent file
+    nonexistent = Path(temp_repo) / "nonexistent.py"
+    chunks = parser.parse_file(nonexistent)
+    assert chunks == []
+
+def test_non_python_files(temp_repo):
+    """Test handling of non-Python files"""
+    parser = CodeParser(temp_repo)
+    
+    # Create a non-Python file
+    non_python = Path(temp_repo) / "readme.md"
+    non_python.write_text("# README")
+    
+    # Should skip non-Python files
+    chunks = parser.parse_file(non_python)
+    assert chunks == []