feat: Extract enums from docstring's type attribute (#30)

* Extract enums from docstring's type attribute. Write tests for enum extraction

* Fix mypy errors and change return type of Parameter.extract_enum method

* style: apply automatic fixes of linters

* Provide JSON serialization for enum sets

* Cover enum type cases suggested by Lars

* Create new class for parameter enums, add more tests for extracting parameter enums from docstring type

* remove backslashes from enum type values

* style: apply automatic fixes of linters

Co-authored-by: Dushko Klincharov <klincarov@halicea.com>
Co-authored-by: duklin <duklin@users.noreply.github.com>
Co-authored-by: Lars Reimann <mail@larsreimann.com>

package_parser/package_parser/commands/get_api/_model.py

@@ -340,6 +340,57 @@ def to_json(self) -> Any:
         }
 
 
+class ParameterEnum:
+    @classmethod
+    def from_docstring(cls, docstring: ParameterAndResultDocstring) -> ParameterEnum:
+        values = set()
+        values.update(ParameterEnum._from_docstring_type(docstring.type))
+        values.update(ParameterEnum._from_docstring_description(docstring.description))
+        return ParameterEnum(values)
+
+    @classmethod
+    def _from_docstring_type(cls, docstring_type: str) -> set[str]:
+        def remove_backslash(e: str):
+            e = e.replace(r"\"", '"')
+            e = e.replace(r"\'", "'")
+            return e
+
+        enum_match = re.search(r"{(.*?)}", docstring_type)
+        values = set()
+        quotes = "'\""
+        if enum_match:
+            enum_str = enum_match.group(1)
+            value = ""
+            inside_value = False
+            curr_quote = None
+            for i, char in enumerate(enum_str):
+                if char in quotes and (i == 0 or (i > 0 and enum_str[i - 1] != "\\")):
+                    if inside_value == False:
+                        inside_value = True
+                        curr_quote = char
+                    elif inside_value == True:
+                        if curr_quote == char:
+                            inside_value = False
+                            curr_quote = None
+                            values.add(remove_backslash(value))
+                            value = ""
+                        else:
+                            value += char
+                elif inside_value:
+                    value += char
+        return values
+
+    @classmethod
+    def _from_docstring_description(cls, docstring_description: str) -> set[str]:
+        return set()
+
+    def __init__(self, values: set[str]) -> None:
+        self.values: set[str] = values
+
+    def to_json(self):
+        return list(self.values)
+
+
 class Parameter:
     @classmethod
     def from_json(cls, json: Any):
@@ -364,6 +415,7 @@ def __init__(
         self.is_public: bool = is_public
         self.assigned_by: ParameterAssignment = assigned_by
         self.docstring = docstring
+        self.enum: ParameterEnum = ParameterEnum.from_docstring(docstring)
 
     def to_json(self) -> Any:
         return {
@@ -372,6 +424,7 @@ def to_json(self) -> Any:
             "is_public": self.is_public,
             "assigned_by": self.assigned_by.name,
             "docstring": self.docstring.to_json(),
+            "enum": self.enum.to_json(),
         }
 
 

package_parser/tests/commands/get_api/test_enums.py

@@ -0,0 +1,38 @@
+from typing import Optional
+
+import pytest
+from package_parser.commands.get_api._model import ParameterEnum
+
+
+@pytest.mark.parametrize(
+    "docstring_type,expected",
+    [
+        ('{"frobenius", "spectral"}, default="frobenius"', {"frobenius", "spectral"}),
+        (
+            "{'strict', 'ignore', 'replace'}, default='strict'",
+            {"strict", "ignore", "replace"},
+        ),
+        (
+            "{'linear', 'poly',             'rbf', 'sigmoid', 'cosine', 'precomputed'}, default='linear'",
+            {"linear", "poly", "rbf", "sigmoid", "cosine", "precomputed"},
+        ),
+        # https://github.com/lars-reimann/sem21/pull/30#discussion_r771288528
+        (r"{\"frobenius\", \'spectral\'}", set()),
+        (r"""{"frobenius'}""", set()),
+        (r"""{'spectral"}""", set()),
+        (r"""{'text\", \"that'}""", {'text", "that'}),
+        (r"""{'text", "that'}""", {'text", "that'}),
+        (r"{'text\', \'that'}", {"text', 'that"}),
+        (r"{'text', 'that'}", {"text", "that"}),
+        (r"""{"text\', \'that"}""", {"text', 'that"}),
+        (r"""{"text', 'that"}""", {"text', 'that"}),
+        (r"""{"text\", \"that"}""", {'text", "that'}),
+        (r'{"text", "that"}', {"text", "that"}),
+        (r"""{\"not', 'be', 'matched'}""", {", ", ", "}),
+        ("""{"gini\\", \\"entropy"}""", {'gini", "entropy'}),
+        ("""{'best\\', \\'random'}""", {"best', 'random"}),
+    ],
+)
+def test_enum_from_docstring_type(docstring_type: str, expected: Optional[set[str]]):
+    result = ParameterEnum._from_docstring_type(docstring_type)
+    assert result == expected