chore: refactor

Author: BobTheBuidler

sphinx_autodoc_vyper/generator.py

@@ -3,7 +3,7 @@
 import os
 from typing import List
 
-from .parser import Constant, Contract, Enum, Event, Function, Struct
+from .parser import Contract, Function, Struct
 
 INDEX_RST = """Vyper Smart Contracts Documentation
 ================================
@@ -74,7 +74,7 @@ def _generate_contract_docs(self, contracts: List[Contract]) -> None:
             if contract.enums:
                 content += _insert_content_section("Enums")
                 for enum in contract.enums:
-                    content += generate_enum_doc(enum)
+                    content += enum.generate_docs()
 
             if contract.structs:
                 content += _insert_content_section("Structs")
@@ -84,12 +84,12 @@ def _generate_contract_docs(self, contracts: List[Contract]) -> None:
             if contract.events:
                 content += _insert_content_section("Events")
                 for event in contract.events:
-                    content += generate_event_doc(event)
+                    content += event.generate_docs()
 
             if contract.constants:
                 content += _insert_content_section("Constants")
                 for constant in contract.constants:
-                    content += generate_constant_doc(constant)
+                    content += constant.generate_docs()
 
             # TODO: fix this
             # if contract.variables:
@@ -137,24 +137,3 @@ def _generate_struct_docs(struct: Struct) -> str:
 def _insert_content_section(name: str) -> str:
     """Insert a hyperlinked content section accessible from the docs index."""
     return f"{name}\n{'-' * len(name)}\n\n"
-
-
-def generate_enum_doc(enum: Enum) -> str:
-    content = f".. py:class:: {enum.name}\n\n"
-    for value in enum.values:
-        content += f"   .. py:attribute:: {value}\n\n"
-    return content
-
-
-def generate_constant_doc(constant: Constant) -> str:
-    return f".. py:data:: {constant.name}\n\n   {constant.type}: {constant.value}\n\n"
-
-
-def generate_event_doc(event: Event) -> str:
-    content = f".. py:class:: {event.name}\n\n"
-    for field in event.fields:
-        content += f"   .. py:attribute:: {field.name}\n\n"
-        content += (
-            f"      {f'indexed({field.type})' if field.indexed else field.type}\n\n"
-        )
-    return content

sphinx_autodoc_vyper/parser.py

@@ -30,6 +30,12 @@ class Enum:
     name: str
     values: List[str]
 
+    def generate_docs(self) -> str:
+        content = f".. py:class:: {self.name}\n\n"
+        for value in self.values:
+            content += f"   .. py:attribute:: {value}\n\n"
+        return content
+
 
 @dataclass
 class Constant:
@@ -43,6 +49,9 @@ def __post_init__(self) -> None:
         if self.type is not None and self.type not in VALID_VYPER_TYPES:
             logger.warning(f"{self} is not a valid Vyper type")
 
+    def generate_docs(self) -> str:
+        return f".. py:data:: {self.name}\n\n   {self.type}: {self.value}\n\n"
+
 
 @dataclass
 class Variable:
@@ -140,6 +149,14 @@ class Event:
     name: str
     params: List[EventParameter]
 
+    def generate_docs(self) -> str:
+        content = f".. py:class:: {event.name}\n\n"
+        for field in event.fields:
+            type_str = f"indexed({field.type})" if field.indexed else field.type
+            content += f"   .. py:attribute:: {field.name}\n\n"
+            content += f"      {type_str}\n\n"
+        return content
+
 
 @dataclass
 class Function:

tests/test_generator.py

@@ -2,12 +2,7 @@
 
 from pathlib import Path
 
-from sphinx_autodoc_vyper.generator import (
-    SphinxGenerator,
-    generate_constant_doc,
-    generate_enum_doc,
-    generate_event_doc,
-)
+from sphinx_autodoc_vyper.generator import SphinxGenerator
 from sphinx_autodoc_vyper.parser import (
     Constant,
     Contract,
@@ -93,17 +88,15 @@ def test_contract_rst_generation(contracts_dir: Path, output_dir: Path) -> None:
 def test_generate_enum_doc() -> None:
     """Test documentation generation for Enum."""
     enum = Enum(name="Status", values=["PENDING", "COMPLETED"])
-    generated_doc = generate_enum_doc(enum)
     expected_doc = ".. py:enum:: Status\n   :members:\n\n   PENDING\n   COMPLETED\n"
-    assert generated_doc == expected_doc
+    assert enum.generate_docs() == expected_doc
 
 
 def test_generate_constant_doc() -> None:
     """Test documentation generation for Constant."""
     constant = Constant(name="MAX_SUPPLY", type="uint256", value="1000000")
-    generated_doc = generate_constant_doc(constant)
     expected_doc = ".. py:data:: MAX_SUPPLY\n\n   :type: uint256\n   :value: 1000000\n"
-    assert generated_doc == expected_doc
+    assert constant.generate_docs() == expected_doc
 
 
 def test_generate_variable_doc() -> None:
@@ -145,14 +138,13 @@ def test_generate_event_doc() -> None:
             Parameter(name="value", type="uint256"),
         ],
     )
-    generated_doc = generate_event_doc(event)
     expected_doc = (
         ".. py:event:: Transfer\n\n"
         "   :param address from: \n"
         "   :param address to: \n"
         "   :param uint256 value: \n"
     )
-    assert generated_doc == expected_doc
+    assert event.generate_docs() == expected_doc
 
 
 def test_generate_docs_for_contract_with_functions_and_structs() -> None: