Skip to content

Commit

Permalink
feat: add kotlin treesitter (#5)
Browse files Browse the repository at this point in the history
fynnfluegge authored Aug 26, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
1 parent ed01c8e commit 4ba33c7
Showing 7 changed files with 103 additions and 3 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
@@ -38,7 +38,7 @@ Focus on writing your code, let AI write the documentation for you. With just a
- [x] Javascript
- [x] Java
- [x] Rust
- [ ] Kotlin
- [x] Kotlin
- [ ] Go
- [ ] C++
- [ ] C
1 change: 1 addition & 0 deletions doc_comments_ai/domain.py
Original file line number Diff line number Diff line change
@@ -3,3 +3,4 @@
from doc_comments_ai.treesitter.treesitter_js import TreesitterJavascript
from doc_comments_ai.treesitter.treesitter_ts import TreesitterTypescript
from doc_comments_ai.treesitter.treesitter_rs import TreesitterRust
from doc_comments_ai.treesitter.treesitter_kt import TreesitterKotlin
47 changes: 47 additions & 0 deletions doc_comments_ai/treesitter/treesitter_kt.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
import tree_sitter
from doc_comments_ai.treesitter.treesitter import (
Treesitter,
TreesitterNode,
get_source_from_node,
)
from doc_comments_ai.constants import Language
from doc_comments_ai.treesitter.treesitter_registry import TreesitterRegistry


class TreesitterKotlin(Treesitter):
def __init__(self):
super().__init__(Language.KOTLIN)

def parse(self, file_bytes: bytes) -> list[TreesitterNode]:
super().parse(file_bytes)
print(self.tree.root_node.sexp())
result = []
methods = self._query_all_methods(self.tree.root_node)
for method in methods:
method_name = self._query_method_name(method["method"])
doc_comment = method["doc_comment"]
result.append(TreesitterNode(method_name, doc_comment, method["method"]))
return result

def _query_method_name(self, node: tree_sitter.Node):
if node.type == "function_declaration":
for child in node.children:
if child.type == "simple_identifier":
return child.text.decode()
return None

def _query_all_methods(self, node: tree_sitter.Node):
methods = []
if node.type == "function_declaration":
doc_comment_node = None
if node.prev_named_sibling and node.prev_named_sibling.type == "comment":
doc_comment_node = get_source_from_node(node.prev_named_sibling)
methods.append({"method": node, "doc_comment": doc_comment_node})
else:
for child in node.children:
methods.extend(self._query_all_methods(child))
return methods


# Register the TreesitterJava class in the registry
TreesitterRegistry.register_treesitter(Language.KOTLIN, TreesitterKotlin)
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "doc-comments-ai"
version = "0.1.4"
version = "0.1.5"
description = ""
authors = ["fynnfluegge <[email protected]>"]
readme = "README.md"
1 change: 1 addition & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
@@ -4,6 +4,7 @@
from tests.fixtures.code_fixture_js import javascript_code_fixture
from tests.fixtures.code_fixture_ts import typescript_code_fixture
from tests.fixtures.code_fixture_rs import rust_code_fixture
from tests.fixtures.code_fixture_kt import kotlin_code_fixture
from tests.fixtures.response_fixtures import (
response_fixture,
response_fixture_language_enclosed,
28 changes: 28 additions & 0 deletions tests/fixtures/code_fixture_kt.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
import pytest


@pytest.fixture
def kotlin_code_fixture():
return """
/**
* Gets 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.
*/
fun getProgrammingLanguage(fileExtension: String): Language {
val languageMapping = HashMap<String, Language>()
languageMapping[".py"] = Language.PYTHON
languageMapping[".js"] = Language.JAVASCRIPT
languageMapping[".ts"] = Language.TYPESCRIPT
languageMapping[".java"] = Language.JAVA
languageMapping[".kt"] = Language.KOTLIN
languageMapping[".lua"] = Language.LUA
return languageMapping.getOrDefault(fileExtension, Language.UNKNOWN)
}
fun getFileExtension(fileName: String): String {
return fileName.substring(fileName.lastIndexOf('.'))
}
"""
25 changes: 24 additions & 1 deletion tests/treesitter_query_test.py
Original file line number Diff line number Diff line change
@@ -141,5 +141,28 @@ def test_rust_qery(rust_code_fixture):
/// ```"""
)

# print(treesitterNodes[1].doc_comment)
assert treesitterNodes[1].doc_comment is None


@pytest.mark.usefixtures("kotlin_code_fixture")
def test_kotlin_qery(kotlin_code_fixture):
tree_sitter = Treesitter.create_treesitter(Language.KOTLIN)
treesitterNodes: list[TreesitterNode] = tree_sitter.parse(
kotlin_code_fixture.encode()
)

assert treesitterNodes.__len__() == 2

assert treesitterNodes[0].name == "getProgrammingLanguage"

assert treesitterNodes[1].name == "getFileExtension"

assert (
treesitterNodes[0].doc_comment
== """/**
* Gets 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.
*/"""
)

0 comments on commit 4ba33c7

Please sign in to comment.