Added support for slice objects in subscriptions (#11981)

Co-authored-by: Bénédikt Tran <[email protected]>

Author: MatrixEditor

CHANGES.rst

@@ -27,6 +27,9 @@ Features added
 * #11892: Improved performance when resolving cross references in cpp domain.
   Patch by Rouslan Korneychuk.
 
+* #11981: Improve rendering of signatures using ``slice`` syntax,
+  e.g., ``def foo(arg: np.float64[:,:]) -> None: ...``.
+
 Bugs fixed
 ----------
 

sphinx/pycode/ast.py

@@ -156,6 +156,20 @@ def visit_Name(self, node: ast.Name) -> str:
     def visit_Set(self, node: ast.Set) -> str:
         return "{" + ", ".join(self.visit(e) for e in node.elts) + "}"
 
+    def visit_Slice(self, node: ast.Slice) -> str:
+        if not node.lower and not node.upper and not node.step:
+            # Empty slice with default values -> [:]
+            return ":"
+
+        start = self.visit(node.lower) if node.lower else ""
+        stop = self.visit(node.upper) if node.upper else ""
+        if not node.step:
+            # Default step size -> [start:stop]
+            return f"{start}:{stop}"
+
+        step = self.visit(node.step) if node.step else ""
+        return f"{start}:{stop}:{step}"
+
     def visit_Subscript(self, node: ast.Subscript) -> str:
         def is_simple_tuple(value: ast.expr) -> bool:
             return (

tests/roots/test-ext-autodoc/target/functions.py

@@ -17,3 +17,6 @@ async def asyncgenerator():
 
 builtin_func = print
 partial_builtin_func = partial(print)
+
+def slice_arg_func(arg: 'float64[:, :]'):
+    pass

tests/test_extensions/test_ext_autodoc_autofunction.py

@@ -199,3 +199,14 @@ def test_async_generator(app):
         '   :async:',
         '',
     ]
+
+
+@pytest.mark.sphinx('html', testroot='ext-autodoc')
+def test_slice_function_arg(app):
+    actual = do_autodoc(app, 'function', 'target.functions.slice_arg_func')
+    assert list(actual) == [
+        '',
+        '.. py:function:: slice_arg_func(arg: float64[:, :])',
+        '   :module: target.functions',
+        '',
+    ]

tests/test_pycode/test_pycode_ast.py

@@ -51,6 +51,13 @@
      "lambda x=0, /, y=1, *args, z, **kwargs: ..."),  # posonlyargs
     ("0x1234", "0x1234"),                       # Constant
     ("1_000_000", "1_000_000"),                 # Constant
+    ("Tuple[:,:]", "Tuple[:, :]"),              # Index, Subscript, 2x Slice
+    ("Tuple[1:2]", "Tuple[1:2]"),               # Index, Subscript, Slice(no-step)
+    ("Tuple[1:2:3]", "Tuple[1:2:3]"),           # Index, Subscript, Slice
+    ("x[:, np.newaxis, :, :]",
+     "x[:, np.newaxis, :, :]"),                 # Index, Subscript, numpy extended syntax
+    ("y[:, 1:3][np.array([0, 2, 4]), :]",
+     "y[:, 1:3][np.array([0, 2, 4]), :]"),       # Index, 2x Subscript, numpy extended syntax
 ])
 def test_unparse(source, expected):
     module = ast.parse(source)