Skip to content

Commit 8ef1486

Browse files
ISOR3Xpawamoy
andauthored
feat: Add method to functions and classes to build and return a stringified signature
Discussion-376: #376 PR-381: #381 Co-authored-by: Timothée Mazzucotelli <[email protected]>
1 parent b346190 commit 8ef1486

File tree

2 files changed

+139
-0
lines changed

2 files changed

+139
-0
lines changed

src/_griffe/models.py

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1654,6 +1654,18 @@ def mro(self) -> list[Class]:
16541654
"""Return a list of classes in order corresponding to Python's MRO."""
16551655
return cast("Class", self.final_target).mro()
16561656

1657+
def signature(self, *, return_type: bool = False, name: str | None = None) -> str:
1658+
"""Construct the class/function signature.
1659+
1660+
Parameters:
1661+
return_type: Whether to include the return type in the signature.
1662+
name: The name of the class/function to use in the signature.
1663+
1664+
Returns:
1665+
A string representation of the class/function signature.
1666+
"""
1667+
return cast("Union[Class, Function]", self.final_target).signature(return_type=return_type, name=name)
1668+
16571669
# SPECIFIC ALIAS METHOD AND PROPERTIES -----------------
16581670
# These methods and properties do not exist on targets,
16591671
# they are specific to aliases.
@@ -1976,6 +1988,23 @@ def parameters(self) -> Parameters:
19761988
except KeyError:
19771989
return Parameters()
19781990

1991+
def signature(self, *, return_type: bool = False, name: str | None = None) -> str:
1992+
"""Construct the class signature.
1993+
1994+
Parameters:
1995+
return_type: Whether to include the return type in the signature.
1996+
name: The name of the class to use in the signature.
1997+
1998+
Returns:
1999+
A string representation of the class signature.
2000+
"""
2001+
all_members = self.all_members
2002+
if "__init__" in all_members:
2003+
init = all_members["__init__"]
2004+
if isinstance(init, Function):
2005+
return init.signature(return_type=return_type, name=name or self.name)
2006+
return ""
2007+
19792008
@property
19802009
def resolved_bases(self) -> list[Object]:
19812010
"""Resolved class bases.
@@ -2125,6 +2154,80 @@ def as_dict(self, **kwargs: Any) -> dict[str, Any]:
21252154
base["returns"] = self.returns
21262155
return base
21272156

2157+
def signature(self, *, return_type: bool = True, name: str | None = None) -> str:
2158+
"""Construct the function signature.
2159+
2160+
Parameters:
2161+
return_type: Whether to include the return type in the signature.
2162+
name: The name of the function to use in the signature.
2163+
2164+
Returns:
2165+
A string representation of the function signature.
2166+
"""
2167+
signature = f"{name or self.name}("
2168+
2169+
has_pos_only = any(p.kind == ParameterKind.positional_only for p in self.parameters)
2170+
render_pos_only_separator = True
2171+
render_kw_only_separator = True
2172+
2173+
param_strs = []
2174+
2175+
for index, param in enumerate(self.parameters):
2176+
# Skip 'self' or 'cls' for class methods if it's the first parameter.
2177+
if index == 0 and param.name in ("self", "cls") and self.parent and self.parent.is_class:
2178+
continue
2179+
2180+
param_str = ""
2181+
2182+
# Handle parameter kind and separators.
2183+
if param.kind != ParameterKind.positional_only:
2184+
if has_pos_only and render_pos_only_separator:
2185+
render_pos_only_separator = False
2186+
param_strs.append("/")
2187+
2188+
if param.kind == ParameterKind.keyword_only and render_kw_only_separator:
2189+
render_kw_only_separator = False
2190+
param_strs.append("*")
2191+
2192+
# Handle variadic parameters.
2193+
if param.kind == ParameterKind.var_positional:
2194+
param_str = "*"
2195+
render_kw_only_separator = False
2196+
elif param.kind == ParameterKind.var_keyword:
2197+
param_str = "**"
2198+
2199+
# Add parameter name.
2200+
param_str += param.name
2201+
2202+
# Handle type annotation
2203+
if param.annotation is not None:
2204+
param_str += f": {param.annotation}"
2205+
equal = " = " # Space around equal when annotation is present.
2206+
else:
2207+
equal = "=" # No space when no annotation.
2208+
2209+
# Handle default value.
2210+
if param.default is not None and param.kind not in {
2211+
ParameterKind.var_positional,
2212+
ParameterKind.var_keyword,
2213+
}:
2214+
param_str += f"{equal}{param.default}"
2215+
2216+
param_strs.append(param_str)
2217+
2218+
# If we have positional-only parameters but no '/' was added yet
2219+
if has_pos_only and render_pos_only_separator:
2220+
param_strs.append("/")
2221+
2222+
signature += ", ".join(param_strs)
2223+
signature += ")"
2224+
2225+
# Add return type if present.
2226+
if return_type and self.annotation:
2227+
signature += f" -> {self.annotation}"
2228+
2229+
return signature
2230+
21282231

21292232
class Attribute(Object):
21302233
"""The class representing a Python module/class/instance attribute."""

tests/test_models.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,14 @@
99

1010
from griffe import (
1111
Attribute,
12+
Class,
1213
Docstring,
14+
Function,
1315
GriffeLoader,
1416
Module,
1517
NameResolutionError,
1618
Parameter,
19+
ParameterKind,
1720
Parameters,
1821
module_vtree,
1922
temporary_inspected_module,
@@ -557,3 +560,36 @@ def __init__(self):
557560
},
558561
) as module:
559562
assert module["A.__init__"].resolve("pd") == "package.mod.pd"
563+
564+
565+
def test_building_function_and_class_signatures() -> None:
566+
"""Test the construction of a class/function signature."""
567+
# Test simple function signatures.
568+
simple_params = Parameters(
569+
Parameter("x", annotation="int"),
570+
Parameter("y", annotation="int", default="0"),
571+
)
572+
simple_func = Function("simple_function", parameters=simple_params, returns="int")
573+
assert simple_func.signature() == "simple_function(x: int, y: int = 0) -> int"
574+
575+
# Test class signatures.
576+
init = Function("__init__", parameters=simple_params, returns="None")
577+
cls = Class("TestClass")
578+
cls.set_member("__init__", init)
579+
assert cls.signature() == "TestClass(x: int, y: int = 0)"
580+
581+
# Create a more complex function with various parameter types.
582+
params = Parameters(
583+
Parameter("a", kind=ParameterKind.positional_only),
584+
Parameter("b", kind=ParameterKind.positional_only, annotation="int", default="0"),
585+
Parameter("c", kind=ParameterKind.positional_or_keyword),
586+
Parameter("d", kind=ParameterKind.positional_or_keyword, annotation="str", default="''"),
587+
Parameter("args", kind=ParameterKind.var_positional),
588+
Parameter("e", kind=ParameterKind.keyword_only),
589+
Parameter("f", kind=ParameterKind.keyword_only, annotation="bool", default="False"),
590+
Parameter("kwargs", kind=ParameterKind.var_keyword),
591+
)
592+
593+
func = Function("test_function", parameters=params, returns="None")
594+
expected = "test_function(a, b: int = 0, /, c, d: str = '', *args, e, f: bool = False, **kwargs) -> None"
595+
assert func.signature() == expected

0 commit comments

Comments
 (0)