Add missing node type of DocString

ducminh-phan · Apr 8, 2019 · 53f80eb · 53f80eb
1 parent 96a1723
commit 53f80eb
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 2 deletions.
diff --git a/main/ast_node/__init__.py b/main/ast_node/__init__.py
@@ -3,6 +3,7 @@
 from .background import Background
 from .comment import Comment
 from .data_table import DataTable
+from .doc_string import DocString
 from .examples import Examples
 from .feature import Feature
 from .gherkin_document import GherkinDocument
@@ -18,6 +19,7 @@
     Background,
     Comment,
     DataTable,
+    DocString,
     Examples,
     Feature,
     GherkinDocument,

diff --git a/main/ast_node/doc_string.py b/main/ast_node/doc_string.py
@@ -0,0 +1,13 @@
+from attr import attrib
+
+from ._base import prepare
+from .location import Location
+
+
+@prepare
+class DocString:
+    location: Location = attrib(cmp=False, repr=False)
+    content: str
+
+    def __iter__(self):
+        yield self
diff --git a/main/ast_node/step.py b/main/ast_node/step.py
@@ -1,9 +1,10 @@
-from typing import Optional
+from typing import Optional, Union
 
 from attr import attrib
 
 from ._base import prepare
 from .data_table import DataTable
+from .doc_string import DocString
 from .location import Location
 
 
@@ -12,7 +13,7 @@ class Step:
     location: Location = attrib(cmp=False, repr=False)
     keyword: str
     text: str
-    argument: Optional[DataTable] = None
+    argument: Optional[Union[DataTable, DocString]] = None
 
     def __iter__(self):
         yield self

diff --git a/main/formatter.py b/main/formatter.py
@@ -7,6 +7,7 @@
     Background,
     Comment,
     DataTable,
+    DocString,
     Examples,
     Feature,
     GherkinDocument,
@@ -167,6 +168,15 @@ def normalize_comment(comment: Comment) -> str:
     return "# " + comment_text
 
 
+def generate_doc_string_lines(docstring: DocString) -> List[str]:
+    raw_lines = docstring.content.splitlines()
+    raw_lines = ['"""'] + raw_lines + ['"""']
+
+    indent_level = INDENT_LEVEL_MAP[Step]
+
+    return [f"{INDENT * indent_level}{line}" for line in raw_lines]
+
+
 ContextMap = Dict[Union[Comment, Tag, TableRow], Any]
 Lines = Iterator[str]
 
@@ -279,3 +289,7 @@ def visit_comment(self, comment: Comment) -> Lines:
 
         comment_text = normalize_comment(comment)
         yield f"{indent}{comment_text}"
+
+    @staticmethod
+    def visit_doc_string(docstring: DocString):
+        return generate_doc_string_lines(docstring)