Implement literal docstring merging for module and class level

Co-authored-by: Josverl <[email protected]>

Author: Copilot

src/stubber/codemod/merge_docstub.py

@@ -242,6 +242,8 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c
         updated_node = self.add_missed_overloads(updated_node, stack_id=())
         # Add any missing @mp_available
         updated_node = self.add_missed_mp_available(updated_node, stack_id=())
+        # Add any missing literal docstrings
+        updated_node = self.add_missed_literal_docstrings(updated_node, stack_id=())
         return updated_node
 
     def add_missed_overloads(self, updated_node: Mod_Class_T, stack_id: tuple) -> Mod_Class_T:
@@ -378,6 +380,83 @@ def add_missed_mp_available(self, updated_node: Mod_Class_T, stack_id: tuple) ->
                 # cst.IndentedBlock(body=tuple(updated_body)))  # type: ignore
         return updated_node
 
+    def add_missed_literal_docstrings(self, updated_node: Mod_Class_T, stack_id: tuple) -> Mod_Class_T:
+        """
+        Add any missing literal docstrings to the updated_node
+        """
+        # Find literal docstrings for this scope
+        literal_docstrings = {}
+        
+        # For module level (empty stack_id), use MODULE_KEY
+        if len(stack_id) == 0:
+            lookup_key = MODULE_KEY
+        else:
+            lookup_key = stack_id
+            
+        if lookup_key in self.annotations and self.annotations[lookup_key].literal_docstrings:
+            literal_docstrings = self.annotations[lookup_key].literal_docstrings
+        
+        if not literal_docstrings:
+            return updated_node
+            
+        # Get the body to modify
+        if isinstance(updated_node, cst.Module):
+            updated_body = list(updated_node.body)
+        elif isinstance(updated_node, cst.ClassDef):
+            updated_body = list(updated_node.body.body)
+        else:
+            raise ValueError(f"Unsupported node type: {updated_node}")
+        
+        # Find assignment statements and insert docstrings after them
+        new_body = []
+        i = 0
+        while i < len(updated_body):
+            stmt = updated_body[i]
+            new_body.append(stmt)
+            
+            # Check if this is an assignment statement for a literal that has a docstring
+            if isinstance(stmt, cst.SimpleStatementLine) and len(stmt.body) == 1:
+                literal_name = None
+                
+                if isinstance(stmt.body[0], cst.Assign):
+                    # Regular assignment: CONST = value
+                    targets = stmt.body[0].targets
+                    if len(targets) == 1 and isinstance(targets[0].target, cst.Name):
+                        literal_name = targets[0].target.value
+                elif isinstance(stmt.body[0], cst.AnnAssign):
+                    # Annotated assignment: CONST: int = value
+                    if isinstance(stmt.body[0].target, cst.Name):
+                        literal_name = stmt.body[0].target.value
+                
+                # If we found a literal with a docstring, check if docstring is missing
+                if literal_name and literal_name in literal_docstrings:
+                    # Check if the next statement is already a docstring for this literal
+                    has_existing_docstring = False
+                    if i + 1 < len(updated_body):
+                        next_stmt = updated_body[i + 1]
+                        if (isinstance(next_stmt, cst.SimpleStatementLine) and 
+                            len(next_stmt.body) == 1 and 
+                            isinstance(next_stmt.body[0], cst.Expr) and
+                            isinstance(next_stmt.body[0].value, cst.SimpleString)):
+                            has_existing_docstring = True
+                    
+                    # If no existing docstring, add the one from doc_stub
+                    if not has_existing_docstring and self.copy_docstr:
+                        docstr_node = literal_docstrings[literal_name]
+                        new_body.append(docstr_node)
+                        log.trace(f"Added literal docstring for {literal_name}")
+            
+            i += 1
+        
+        # Update the node with the new body
+        if isinstance(updated_node, cst.Module):
+            updated_node = updated_node.with_changes(body=tuple(new_body))
+        elif isinstance(updated_node, cst.ClassDef):
+            new_indented_body = updated_node.body.with_changes(body=tuple(new_body))
+            updated_node = updated_node.with_changes(body=new_indented_body)
+        
+        return updated_node
+
     def locate_function_by_name(self, overload, updated_body):
         """locate the (last) function by name"""
         matched = False
@@ -426,6 +505,8 @@ def leave_ClassDef(self, original_node: cst.ClassDef, updated_node: cst.ClassDef
         updated_node = self.add_missed_overloads(updated_node, stack_id)
         # Add any missing @mp_available
         updated_node = self.add_missed_mp_available(updated_node, stack_id)
+        # Add any missing literal docstrings
+        updated_node = self.add_missed_literal_docstrings(updated_node, stack_id)
         return updated_node
 
     # ------------------------------------------------------------------------

src/stubber/typing_collector.py

@@ -33,6 +33,8 @@ class AnnoValue:
     "function / method overloads read from the docstub source"
     mp_available: List[TypeInfo] = field(default_factory=list)
     "function / method `overloads` read from the docstub source"
+    literal_docstrings: Dict[str, cst.SimpleStatementLine] = field(default_factory=dict)
+    "literal/constant name -> docstring node mappings for literal docstrings"
 
 
 class TransformError(Exception):
@@ -85,12 +87,58 @@ def __init__(self):
         ] = {}
         self.comments: List[str] = []
 
+    def _collect_literal_docstrings(self, body: Sequence[cst.BaseStatement]) -> Dict[str, cst.SimpleStatementLine]:
+        """
+        Collect literal docstrings from a sequence of statements.
+        Looks for patterns like:
+        CONSTANT = value
+        '''docstring for constant'''
+        """
+        literal_docstrings = {}
+        
+        for i, stmt in enumerate(body):
+            # Look for assignment statements  
+            if isinstance(stmt, cst.SimpleStatementLine) and len(stmt.body) == 1:
+                if isinstance(stmt.body[0], (cst.Assign, cst.AnnAssign)):
+                    # Get the literal name
+                    literal_name = None
+                    if isinstance(stmt.body[0], cst.Assign):
+                        # Handle regular assignment: CONST = value
+                        targets = stmt.body[0].targets
+                        if len(targets) == 1 and isinstance(targets[0].target, cst.Name):
+                            literal_name = targets[0].target.value
+                    elif isinstance(stmt.body[0], cst.AnnAssign):
+                        # Handle annotated assignment: CONST: int = value
+                        if isinstance(stmt.body[0].target, cst.Name):
+                            literal_name = stmt.body[0].target.value
+                    
+                    # Check if the next statement is a docstring
+                    if literal_name and i + 1 < len(body):
+                        next_stmt = body[i + 1]
+                        if (isinstance(next_stmt, cst.SimpleStatementLine) and 
+                            len(next_stmt.body) == 1 and 
+                            isinstance(next_stmt.body[0], cst.Expr) and
+                            isinstance(next_stmt.body[0].value, cst.SimpleString)):
+                            # Found a literal with a following docstring
+                            literal_docstrings[literal_name] = next_stmt
+                            
+        return literal_docstrings
+
     # ------------------------------------------------------------
     def visit_Module(self, node: cst.Module) -> bool:
-        """Store the module docstring"""
+        """Store the module docstring and collect literal docstrings"""
+        # Store module docstring
         docstr = node.get_docstring()
         if docstr:
             self.annotations[MODULE_KEY] = AnnoValue(docstring=docstr)
+        
+        # Collect module-level literal docstrings
+        literal_docstrings = self._collect_literal_docstrings(node.body)
+        if literal_docstrings:
+            if MODULE_KEY not in self.annotations:
+                self.annotations[MODULE_KEY] = AnnoValue()
+            self.annotations[MODULE_KEY].literal_docstrings.update(literal_docstrings)
+            
         return True
 
     def visit_Comment(self, node: cst.Comment) -> None:
@@ -108,6 +156,7 @@ def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]:
         """
         collect info from a classdef:
         - name, decorators, docstring
+        - class-level literal docstrings
         """
         # "Store the class docstring
         docstr_node = self.update_append_first_node(node)
@@ -120,7 +169,17 @@ def visit_ClassDef(self, node: cst.ClassDef) -> Optional[bool]:
             def_type="classdef",
             def_node=node,
         )
-        self.annotations[tuple(self.stack)] = AnnoValue(type_info=ti)
+        
+        # Collect class-level literal docstrings
+        literal_docstrings = {}
+        if isinstance(node.body, cst.IndentedBlock):
+            literal_docstrings = self._collect_literal_docstrings(node.body.body)
+        
+        anno_value = AnnoValue(type_info=ti)
+        if literal_docstrings:
+            anno_value.literal_docstrings.update(literal_docstrings)
+            
+        self.annotations[tuple(self.stack)] = anno_value
 
     def leave_ClassDef(self, original_node: cst.ClassDef) -> None:
         """remove the class name from the stack"""