8
8
9
9
import contextlib
10
10
import io
11
+ import keyword
11
12
import re
12
13
import tokenize
13
14
from typing import Any , Final , MutableMapping , MutableSequence , NamedTuple , Sequence , Tuple
@@ -35,12 +36,16 @@ class ArgSig:
35
36
36
37
def __init__ (self , name : str , type : str | None = None , default : bool = False ):
37
38
self .name = name
38
- if type and not is_valid_type (type ):
39
- raise ValueError ("Invalid type: " + type )
40
39
self .type = type
41
40
# Does this argument have a default value?
42
41
self .default = default
43
42
43
+ def is_star_arg (self ) -> bool :
44
+ return self .name .startswith ("*" ) and not self .name .startswith ("**" )
45
+
46
+ def is_star_kwarg (self ) -> bool :
47
+ return self .name .startswith ("**" )
48
+
44
49
def __repr__ (self ) -> str :
45
50
return "ArgSig(name={}, type={}, default={})" .format (
46
51
repr (self .name ), repr (self .type ), repr (self .default )
@@ -59,7 +64,68 @@ def __eq__(self, other: Any) -> bool:
59
64
class FunctionSig (NamedTuple ):
60
65
name : str
61
66
args : list [ArgSig ]
62
- ret_type : str
67
+ ret_type : str | None
68
+
69
+ def is_special_method (self ) -> bool :
70
+ return bool (
71
+ self .name .startswith ("__" )
72
+ and self .name .endswith ("__" )
73
+ and self .args
74
+ and self .args [0 ].name in ("self" , "cls" )
75
+ )
76
+
77
+ def has_catchall_args (self ) -> bool :
78
+ """Return if this signature has catchall args: (*args, **kwargs)"""
79
+ if self .args and self .args [0 ].name in ("self" , "cls" ):
80
+ args = self .args [1 :]
81
+ else :
82
+ args = self .args
83
+ return (
84
+ len (args ) == 2
85
+ and all (a .type in (None , "Any" , "typing.Any" ) for a in args )
86
+ and args [0 ].is_star_arg ()
87
+ and args [1 ].is_star_kwarg ()
88
+ )
89
+
90
+ def is_identity (self ) -> bool :
91
+ """Return if this signature is the catchall identity: (*args, **kwargs) -> Any"""
92
+ return self .has_catchall_args () and self .ret_type in (None , "Any" , "typing.Any" )
93
+
94
+ def format_sig (self , any_val : str | None = None , suffix : str = ": ..." ) -> str :
95
+ args : list [str ] = []
96
+ for arg in self .args :
97
+ arg_def = arg .name
98
+
99
+ if arg_def in keyword .kwlist :
100
+ arg_def = "_" + arg_def
101
+
102
+ if (
103
+ arg .type is None
104
+ and any_val is not None
105
+ and arg .name not in ("self" , "cls" )
106
+ and not arg .name .startswith ("*" )
107
+ ):
108
+ arg_type : str | None = any_val
109
+ else :
110
+ arg_type = arg .type
111
+ if arg_type :
112
+ arg_def += ": " + arg_type
113
+ if arg .default :
114
+ arg_def += " = ..."
115
+
116
+ elif arg .default :
117
+ arg_def += "=..."
118
+
119
+ args .append (arg_def )
120
+
121
+ retfield = ""
122
+ ret_type = self .ret_type if self .ret_type else any_val
123
+ if ret_type is not None :
124
+ retfield = " -> " + ret_type
125
+
126
+ return "def {name}({args}){ret}{suffix}" .format (
127
+ name = self .name , args = ", " .join (args ), ret = retfield , suffix = suffix
128
+ )
63
129
64
130
65
131
# States of the docstring parser.
@@ -176,17 +242,17 @@ def add_token(self, token: tokenize.TokenInfo) -> None:
176
242
177
243
# arg_name is empty when there are no args. e.g. func()
178
244
if self .arg_name :
179
- try :
245
+ if self .arg_type and not is_valid_type (self .arg_type ):
246
+ # wrong type, use Any
247
+ self .args .append (
248
+ ArgSig (name = self .arg_name , type = None , default = bool (self .arg_default ))
249
+ )
250
+ else :
180
251
self .args .append (
181
252
ArgSig (
182
253
name = self .arg_name , type = self .arg_type , default = bool (self .arg_default )
183
254
)
184
255
)
185
- except ValueError :
186
- # wrong type, use Any
187
- self .args .append (
188
- ArgSig (name = self .arg_name , type = None , default = bool (self .arg_default ))
189
- )
190
256
self .arg_name = ""
191
257
self .arg_type = None
192
258
self .arg_default = None
@@ -240,7 +306,7 @@ def args_kwargs(signature: FunctionSig) -> bool:
240
306
241
307
242
308
def infer_sig_from_docstring (docstr : str | None , name : str ) -> list [FunctionSig ] | None :
243
- """Convert function signature to list of TypedFunctionSig
309
+ """Convert function signature to list of FunctionSig
244
310
245
311
Look for function signatures of function in docstring. Signature is a string of
246
312
the format <function_name>(<signature>) -> <return type> or perhaps without
0 commit comments