Merge pull request #660 from kkoomen/feature/python-classes

Implement class docstrings
kkoomen · May 1, 2024 · ca6f40e · ca6f40e
2 parents 39b42ec + 40422f2
commit ca6f40e
Show file tree

Hide file tree

Showing 11 changed files with 234 additions and 0 deletions.
diff --git a/helper/src/python/docs/doxygen.yaml b/helper/src/python/docs/doxygen.yaml
@@ -1,6 +1,12 @@
 # https://www.woolseyworkshop.com/2020/06/25/documenting-python-programs-with-doxygen/
 
 templates:
+  class:
+    node_types:
+      - class_definition
+    template: |
+      """! @brief [TODO:description]"""
+
   function:
     node_types:
       - function_definition

diff --git a/helper/src/python/docs/google.yaml b/helper/src/python/docs/google.yaml
@@ -1,6 +1,21 @@
 # https://github.com/google/styleguide/blob/gh-pages/pyguide.md
 
 templates:
+  class:
+    node_types:
+      - class_definition
+    template: |
+      """
+      [TODO:description]
+      {% if attributes %}
+      ~
+      Attributes:
+        {% for attribute in attributes %}
+      <INDENT>{{ attribute.name }}{% if show_types or not attribute.type %} ({{ attribute.type | default(value="[TODO:type]") }}){% endif %}: [TODO:description]
+        {% endfor %}
+      {% endif %}
+      """
+
   function:
     node_types:
       - function_definition

diff --git a/helper/src/python/docs/numpy.yaml b/helper/src/python/docs/numpy.yaml
@@ -1,6 +1,23 @@
 # http://daouzli.com/blog/docstring.html#numpydoc
 
 templates:
+  class:
+    node_types:
+      - class_definition
+    template: |
+      """
+      [TODO:description]
+      {% if attributes %}
+      ~
+      Attributes
+      ----------
+        {% for attribute in attributes %}
+      {{ attribute.name }}{% if show_types or not attribute.type %} : {{ attribute.type | default(value="[TODO:type]") }}{% endif %}
+      <INDENT>[TODO:description]
+        {% endfor %}
+      {% endif %}
+      """
+
   function:
     node_types:
       - function_definition

diff --git a/helper/src/python/docs/reST.yaml b/helper/src/python/docs/reST.yaml
@@ -1,6 +1,21 @@
 # http://daouzli.com/blog/docstring.html#restructuredtext
 
 templates:
+  class:
+    node_types:
+      - class_definition
+    template: |
+      """
+      [TODO:description]
+      {% if attributes %}
+      ~
+      Attributes:
+        {% for attribute in attributes %}
+      <INDENT>{{ attribute.name }}{% if show_types or not attribute.type %} ({{ attribute.type | default(value="[TODO:type]") }}){% endif %}: [TODO:description]
+        {% endfor %}
+      {% endif %}
+      """
+
   function:
     node_types:
       - function_definition

diff --git a/helper/src/python/docs/sphinx.yaml b/helper/src/python/docs/sphinx.yaml
@@ -1,6 +1,21 @@
 # https://sphinx-rtd-tutorial.readthedocs.io/en/latest/docstrings.html
 
 templates:
+  class:
+    node_types:
+      - class_definition
+    template: |
+      """
+      [TODO:description]
+      {% if attributes %}
+      ~
+      Attributes:
+        {% for attribute in attributes %}
+      <INDENT>{{ attribute.name }}{% if show_types or not attribute.type %} ({{ attribute.type | default(value="[TODO:type]") }}){% endif %}: [TODO:description]
+        {% endfor %}
+      {% endif %}
+      """
+
   function:
     node_types:
       - function_definition

diff --git a/helper/src/python/parser.rs b/helper/src/python/parser.rs
@@ -64,6 +64,7 @@ impl<'a> PythonParser<'a> {
         for child_node in traverse::PreOrder::new(node.walk()) {
             if child_node.start_position().row + 1 == *self.line && self.node_types.contains(&child_node.kind()) {
                 return match child_node.kind() {
+                    "class_definition" => Some(self.parse_class(&child_node)),
                     "function_definition" => Some(self.parse_function(&child_node)),
                     _ => None,
                 };
@@ -73,6 +74,60 @@ impl<'a> PythonParser<'a> {
         None
     }
 
+    fn parse_class(&self, node: &Node) -> Result<Map<String, Value>, String> {
+        let mut tokens = Map::new();
+
+        for child_node in node.children(&mut node.walk()) {
+            match child_node.kind() {
+                "identifier" => {
+                    tokens.insert("name".to_string(), Value::String(self.get_node_text(&child_node)));
+                },
+                "block" => {
+                    let attributes = self.parse_class_attributes(&child_node);
+                    if !attributes.is_empty()  {
+                        tokens.insert("attributes".to_string(), Value::Array(attributes));
+                    }
+                }
+                _ => {},
+            }
+        }
+
+        Ok(tokens)
+    }
+
+    fn parse_class_attributes(&self, node: &Node) -> Vec<Value> {
+        let mut attributes = Vec::new();
+
+        node
+            .children(&mut node.walk())
+            .filter(|node| node.kind() == "expression_statement")
+            .for_each(|node| {
+                let mut attr = Map::new();
+
+                node
+                    .children(&mut node.walk())
+                    .filter(|node| node.kind() == "assignment")
+                    .for_each(|node|
+                        for child_node in node.children(&mut node.walk()) {
+                            match child_node.kind() {
+                                "identifier" => {
+                                    attr.insert("name".to_string(), Value::String(self.get_node_text(&child_node)));
+                                },
+                                "type" => {
+                                    attr.insert("type".to_string(), Value::String(self.get_node_text(&child_node)));
+                                },
+                                _ => {}
+                            }
+                        }
+                    );
+
+                attributes.push(Value::Object(attr));
+            });
+
+        attributes
+    }
+
+
     fn parse_function(&self, node: &Node) -> Result<Map<String, Value>, String> {
         let mut tokens = Map::new();
 

diff --git a/test/filetypes/python/classes-doc-doxygen.vader b/test/filetypes/python/classes-doc-doxygen.vader
@@ -0,0 +1,17 @@
+# ==============================================================================
+# Classes with attributes
+# ==============================================================================
+Given python (class with two attributes):
+  class Foo:
+    bar = 1
+    baz: int = 2
+
+Do (run doge):
+  :let b:doge_doc_standard='doxygen'\<CR>
+  \<C-d>
+
+Expect python (class with docstring describing all attributes):
+  class Foo:
+    """! @brief [TODO:description]"""
+    bar = 1
+    baz: int = 2
diff --git a/test/filetypes/python/classes-doc-google.vader b/test/filetypes/python/classes-doc-google.vader
@@ -0,0 +1,23 @@
+# ==============================================================================
+# Classes with attributes
+# ==============================================================================
+Given python (class with two attributes):
+  class Foo:
+    bar = 1
+    baz: int = 2
+
+Do (run doge):
+  :let b:doge_doc_standard='google'\<CR>
+  \<C-d>
+
+Expect python (class with docstring describing all attributes):
+  class Foo:
+    """
+    [TODO:description]
+
+    Attributes:
+      bar ([TODO:type]): [TODO:description]
+      baz: [TODO:description]
+    """
+    bar = 1
+    baz: int = 2
diff --git a/test/filetypes/python/classes-doc-numpy.vader b/test/filetypes/python/classes-doc-numpy.vader
@@ -0,0 +1,26 @@
+# ==============================================================================
+# Classes with attributes
+# ==============================================================================
+Given python (class with two attributes):
+  class Foo:
+    bar = 1
+    baz: int = 2
+
+Do (run doge):
+  :let b:doge_doc_standard='numpy'\<CR>
+  \<C-d>
+
+Expect python (class with docstring describing all attributes):
+  class Foo:
+    """
+    [TODO:description]
+
+    Attributes
+    ----------
+    bar : [TODO:type]
+      [TODO:description]
+    baz
+      [TODO:description]
+    """
+    bar = 1
+    baz: int = 2
diff --git a/test/filetypes/python/classes-doc-sphinx.vader b/test/filetypes/python/classes-doc-sphinx.vader
@@ -0,0 +1,23 @@
+# ==============================================================================
+# Classes with attributes
+# ==============================================================================
+Given python (class with two attributes):
+  class Foo:
+    bar = 1
+    baz: int = 2
+
+Do (run doge):
+  :let b:doge_doc_standard='sphinx'\<CR>
+  \<C-d>
+
+Expect python (class with docstring describing all attributes):
+  class Foo:
+    """
+    [TODO:description]
+
+    Attributes:
+      bar ([TODO:type]): [TODO:description]
+      baz: [TODO:description]
+    """
+    bar = 1
+    baz: int = 2
diff --git a/test/filetypes/python/classes.vader b/test/filetypes/python/classes.vader
@@ -0,0 +1,22 @@
+# ==============================================================================
+# Classes with attributes
+# ==============================================================================
+Given python (class with two attributes):
+  class Foo:
+    bar = 1
+    baz: int = 2
+
+Do (run doge):
+  \<C-d>
+
+Expect python (class with docstring describing all attributes):
+  class Foo:
+    """
+    [TODO:description]
+
+    Attributes:
+      bar ([TODO:type]): [TODO:description]
+      baz: [TODO:description]
+    """
+    bar = 1
+    baz: int = 2