extract_descriptions.py should support test functions as well

pyproject.toml

@@ -13,15 +13,14 @@ readme = "README.pypi.md"
 version = "2.4.0"
 
 [tool.poetry.dependencies]
-python = ">=3.8.1,<3.12"
 aiohttp = {extras = ["speedups"], version = "*"}
 arch = "*"
 bert-score = ">=0.3.13"
 catboost = "*"
 evaluate = "*"
 ipywidgets = "*"
 kaleido = ">=0.2.1,!=0.2.1.post1"
-langchain-openai = { version = ">=0.1.8", optional = true }
+langchain-openai = {version = ">=0.1.8", optional = true}
 langdetect = "*"
 latex2mathml = ">=3.77.0"
 llvmlite = {version = "*", python = ">=3.8,<=3.11"}
@@ -36,6 +35,7 @@ plotly = "*"
 plotly-express = "*"
 polars = "*"
 pycocoevalcap = {version = "^1.2", optional = true}
+python = ">=3.8.1,<3.12"
 python-dotenv = "*"
 ragas = {version = ">=0.1.7", optional = true}
 rouge = ">=1"
@@ -82,6 +82,7 @@ all = [
   "ragas",
   "langchain-openai",
 ]
+huggingface = ["transformers"]
 llm = [
   "torch",
   "transformers",
@@ -92,7 +93,6 @@ llm = [
 ]
 pytorch = ["torch"]
 r-support = ["rpy2"]
-huggingface = ["transformers"]
 
 [build-system]
 build-backend = "poetry.core.masonry.api"

scripts/extract_descriptions.py

@@ -6,6 +6,7 @@
 
  - path: path to a test file or directory containing test files
 """
+
 import os
 import shutil
 
@@ -15,6 +16,15 @@
 output_dir_path = "build/_test_descriptions"
 
 
+def is_test_function_signature(line, previous_line):
+    """
+    Test functions should have a @tags or @tasks decorator call on top of them
+    """
+    return line.startswith("def") and (
+        "@tags" in previous_line or "@tasks" in previous_line
+    )
+
+
 def retrieve_test_description(path):
     """Generate a test description using gpt4
     You can switch to gpt3.5 if you don't have access but gpt4 should do a better job
@@ -25,8 +35,26 @@ def retrieve_test_description(path):
     # the description should be inserted after the class definition line
     existing_description_lines = []
     lines = file_contents.split("\n")
+
+    advance_to_next_line = False
     for i, line in enumerate(lines):
-        if line.startswith("class"):
+        previous_line = lines[i - 1] if i > 0 else ""
+
+        if advance_to_next_line or (
+            line.startswith("class") or is_test_function_signature(line, previous_line)
+        ):
+            # ensure this is not a multi-line function signature like this:
+            #
+            # def test_function(
+            #     arg1,
+            #     arg2
+            # ):
+            #
+            # we want to keep iterating until we find the closing parenthesis
+            if ")" not in line:
+                advance_to_next_line = True
+                continue
+
             # check if there is already a doc string for the class
             if '"""' in lines[i + 1]:
                 existing_description_lines.append(i + 1)
@@ -36,6 +64,8 @@ def retrieve_test_description(path):
                     if '"""' in lines[j]:
                         break
                     j += 1
+
+            advance_to_next_line = False
             break
 
     existing_description_lines = [