Added support for Haskell Language (#28)

Added support for Haskell language.

---------

Co-authored-by: aravind.mallapureddy <[email protected]>
Co-authored-by: fynnfluegge <[email protected]>

README.md

@@ -97,6 +97,7 @@ aicomment <RELATIVE_FILE_PATH> --guided
 - [x] C++
 - [x] C
 - [x] C#
+- [x] Haskell
 
 ## 📋 Requirements
 

doc_comments_ai/constants.py

@@ -22,4 +22,5 @@ class Language(Enum):
     PERL = "perl"
     LUA = "lua"
     R = "r"
+    HASKELL = "haskell"
     UNKNOWN = "unknown"

doc_comments_ai/treesitter/__init__.py

@@ -10,3 +10,4 @@
 from doc_comments_ai.treesitter.treesitter_py import TreesitterPython
 from doc_comments_ai.treesitter.treesitter_rs import TreesitterRust
 from doc_comments_ai.treesitter.treesitter_ts import TreesitterTypescript
+from doc_comments_ai.treesitter.treesitter_hs import TreesitterHaskell

doc_comments_ai/treesitter/treesitter.py

@@ -62,6 +62,15 @@ def _query_all_methods(
                 and node.prev_named_sibling.type == self.doc_comment_identifier
             ):
                 doc_comment_node = node.prev_named_sibling.text.decode()
+            else:
+                # added for haskell purpose.
+                if node.prev_named_sibling.type == "signature":
+                    prev_node = node.prev_named_sibling
+                    if (
+                        prev_node.prev_named_sibling
+                        and prev_node.prev_named_sibling.type == self.doc_comment_identifier
+                    ):
+                        doc_comment_node = prev_node.prev_named_sibling.text.decode()
             methods.append({"method": node, "doc_comment": doc_comment_node})
         else:
             for child in node.children:

doc_comments_ai/treesitter/treesitter_hs.py

@@ -0,0 +1,23 @@
+import tree_sitter
+
+from doc_comments_ai.constants import Language
+from doc_comments_ai.treesitter.treesitter import Treesitter
+from doc_comments_ai.treesitter.treesitter_registry import TreesitterRegistry
+
+
+class TreesitterHaskell(Treesitter):
+    def __init__(self):
+        super().__init__(
+            Language.HASKELL, "function", "variable", "comment"
+        )
+
+    def _query_method_name(self, node: tree_sitter.Node):
+        if node.type == self.method_declaration_identifier:
+            for child in node.children:
+                if child.type == self.method_name_identifier:
+                    return child.text.decode()
+        return None
+
+
+# Register the TreesitterHaskell class in the registry
+TreesitterRegistry.register_treesitter(Language.HASKELL, TreesitterHaskell)

doc_comments_ai/utils.py

@@ -34,6 +34,7 @@ def get_programming_language(file_extension: str) -> Language:
         ".cpp": Language.CPP,
         ".c": Language.C,
         ".cs": Language.C_SHARP,
+        ".hs": Language.HASKELL,
     }
     return language_mapping.get(file_extension, Language.UNKNOWN)
 

pyproject.toml

@@ -22,6 +22,9 @@ inquirer = "^3.1.3"
 [tool.poetry.scripts]
 aicomment = "doc_comments_ai.__main__:main"
 
+[tool.poetry.group.dev.dependencies]
+pytest = "^7.4.4"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

tests/conftest.py

@@ -10,6 +10,7 @@
 from tests.fixtures.code_fixture_py import python_code_fixture
 from tests.fixtures.code_fixture_rs import rust_code_fixture
 from tests.fixtures.code_fixture_ts import typescript_code_fixture
+from tests.fixtures.code_fixture_hs import haskell_code_fixture
 from tests.fixtures.response_fixtures import (
     response_fixture, response_fixture_language_enclosed,
     response_fixture_with_text)

tests/fixtures/code_fixture_cs.py

@@ -35,10 +35,10 @@ def csharp_code_fixture():
         languageMapping[".kt"] = Language.KOTLIN;
         languageMapping[".lua"] = Language.LUA;
 
-        // Return the corresponding language if it exists in the mapping, otherwise return Language.UNKNOWN 
+        // Return the corresponding language if it exists in the mapping, otherwise return Language.UNKNOWN
         return languageMapping.TryGetValue(fileExtension, out var language) ? language : Language.UNKNOWN;
     }
-    
+
     public static string GetFileExtension(string fileName)
     {
         int dotPosition = fileName.LastIndexOf('.');

tests/fixtures/code_fixture_hs.py

@@ -0,0 +1,36 @@
+import pytest
+
+
+@pytest.fixture
+def haskell_code_fixture():
+    return """
+module Test where
+
+import Prelude
+import qualified Data.HashMap.String as HM
+
+-- Enumeration representing various programming languages.
+data Language = PYTHON | JAVASCRIPT | TYPESCRIPT | JAVA | KOTLIN | LUA | UNKNOWN
+
+{-
+  Get the corresponding programming language based on the given file extension.
+
+  @param fileExtension The file extension of the programming file.
+  @return The corresponding programming language if it exists in the mapping, otherwise Language.UNKNOWN.
+-}
+getProgrammingLanguage :: String -> Language
+getProgrammingLanguage fileExtension =
+    let languageMapping = HM.insert ".py" PYTHON
+                          $ HM.insert ".js" JAVASCRIPT
+                          $ HM.insert ".ts" TYPESCRIPT
+                          $ HM.insert ".java" JAVA
+                          $ HM.insert ".kt" KOTLIN
+                          $ HM.singleton ".lua" LUA
+    in case lookup fileExtension languageMapping of
+         Just v -> v
+         Nothing -> UNKNOWN
+getFileExtension :: String -> String
+getFileExtension fileName =
+    let dot = dropWhile ((\=) '.') fileName
+    in dot
+"""

tests/fixtures/code_fixture_java.py

@@ -34,7 +34,7 @@ def java_code_fixture():
         languageMapping.put(".kt", Language.KOTLIN);
         languageMapping.put(".lua", Language.LUA);
 
-        // Return the corresponding language if it exists in the mapping, otherwise return Language.UNKNOWN 
+        // Return the corresponding language if it exists in the mapping, otherwise return Language.UNKNOWN
         return languageMapping.getOrDefault(fileExtension, Language.UNKNOWN);
     }
 

tests/treesitter_query_test.py

@@ -93,7 +93,7 @@ def test_java_query(java_code_fixture):
         languageMapping.put(".kt", Language.KOTLIN);
         languageMapping.put(".lua", Language.LUA);
 
-        // Return the corresponding language if it exists in the mapping, otherwise return Language.UNKNOWN 
+        // Return the corresponding language if it exists in the mapping, otherwise return Language.UNKNOWN
         return languageMapping.getOrDefault(fileExtension, Language.UNKNOWN);
     }"""
     )
@@ -465,7 +465,46 @@ def test_csharp_query(csharp_code_fixture):
         languageMapping[".kt"] = Language.KOTLIN;
         languageMapping[".lua"] = Language.LUA;
 
-        // Return the corresponding language if it exists in the mapping, otherwise return Language.UNKNOWN 
+        // Return the corresponding language if it exists in the mapping, otherwise return Language.UNKNOWN
         return languageMapping.TryGetValue(fileExtension, out var language) ? language : Language.UNKNOWN;
     }"""
     )
+
+@pytest.mark.usefixtures("haskell_code_fixture")
+def test_hs_query(haskell_code_fixture):
+    tree_sitter = Treesitter.create_treesitter(Language.HASKELL)
+    treesitterNodes: list[TreesitterMethodNode] = tree_sitter.parse(
+        haskell_code_fixture.encode()
+    )
+
+    assert treesitterNodes.__len__() == 2
+
+    assert treesitterNodes[0].name == "getProgrammingLanguage"
+
+    assert treesitterNodes[1].name == "getFileExtension"
+
+    assert (
+        treesitterNodes[0].doc_comment
+        == """{-
+  Get the corresponding programming language based on the given file extension.
+
+  @param fileExtension The file extension of the programming file.
+  @return The corresponding programming language if it exists in the mapping, otherwise Language.UNKNOWN.
+-}"""
+    )
+
+    assert treesitterNodes[1].doc_comment is None
+
+    assert (
+        treesitterNodes[0].method_source_code
+        == """getProgrammingLanguage fileExtension =
+    let languageMapping = HM.insert ".py" PYTHON
+                          $ HM.insert ".js" JAVASCRIPT
+                          $ HM.insert ".ts" TYPESCRIPT
+                          $ HM.insert ".java" JAVA
+                          $ HM.insert ".kt" KOTLIN
+                          $ HM.singleton ".lua" LUA
+    in case lookup fileExtension languageMapping of
+         Just v -> v
+         Nothing -> UNKNOWN"""
+    )