File: infer.py

package info (click to toggle)
pyannotate 1.2.0-2
  • links: PTS, VCS
  • area: main
  • in suites: bookworm, bullseye
  • size: 384 kB
  • sloc: python: 5,001; makefile: 5
file content (243 lines) | stat: -rw-r--r-- 8,522 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
"""Infer an annotation from a set of concrete runtime type signatures.

The main entry point is 'infer_annotation'.
"""

from typing import Dict, Iterable, List, Optional, Set, Tuple

from pyannotate_tools.annotations.parse import parse_type_comment
from pyannotate_tools.annotations.types import (
    AbstractType,
    AnyType,
    ARG_POS,
    Argument,
    ClassType,
    is_optional,
    TupleType,
    UnionType,
    NoReturnType,
)

IGNORED_ITEMS = {
    'unittest.mock.Mock',
    'unittest.mock.MagicMock',
    'mock.mock.Mock',
    'mock.mock.MagicMock',
}

class InferError(Exception):
    """Raised if we can't infer a signature for some reason."""


def infer_annotation(type_comments):
    # type: (List[str]) -> Tuple[List[Argument], AbstractType]
    """Given some type comments, return a single inferred signature.

    Args:
        type_comments: Strings of form '(arg1, ... argN) -> ret'

    Returns: Tuple of (argument types and kinds, return type).
    """
    assert type_comments
    args = {}  # type: Dict[int, Set[Argument]]
    returns = set()
    for comment in type_comments:
        arg_types, return_type = parse_type_comment(comment)
        for i, arg_type in enumerate(arg_types):
            args.setdefault(i, set()).add(arg_type)
        returns.add(return_type)
    combined_args = []
    for i in sorted(args):
        arg_infos = list(args[i])
        kind = argument_kind(arg_infos)
        if kind is None:
            raise InferError('Ambiguous argument kinds:\n' + '\n'.join(type_comments))
        types = [arg.type for arg in arg_infos]
        combined = combine_types(types)
        if str(combined) == 'None':
            # It's very rare for an argument to actually be typed `None`, more likely than
            # not we simply don't have any data points for this argument.
            combined = UnionType([ClassType('None'), AnyType()])
        if kind != ARG_POS and (len(str(combined)) > 120 or isinstance(combined, UnionType)):
            # Avoid some noise.
            combined = AnyType()
        combined_args.append(Argument(combined, kind))
    combined_return = combine_types(returns)
    return combined_args, combined_return


def argument_kind(args):
    # type: (List[Argument]) -> Optional[str]
    """Return the kind of an argument, based on one or more descriptions of the argument.

    Return None if every item does not have the same kind.
    """
    kinds = set(arg.kind for arg in args)
    if len(kinds) != 1:
        return None
    return kinds.pop()


def combine_types(types):
    # type: (Iterable[AbstractType]) -> AbstractType
    """Given some types, return a combined and simplified type.

    For example, if given 'int' and 'List[int]', return Union[int, List[int]]. If given
    'int' and 'int', return just 'int'.
    """
    items = simplify_types(types)
    if len(items) == 1:
        return items[0]
    else:
        return UnionType(items)


def simplify_types(types):
    # type: (Iterable[AbstractType]) -> List[AbstractType]
    """Given some types, give simplified types representing the union of types."""
    flattened = flatten_types(types)
    items = filter_ignored_items(flattened)
    items = [simplify_recursive(item) for item in items]
    items = merge_items(items)
    items = dedupe_types(items)
    # We have to remove reundant items after everything has been simplified and
    # merged as this simplification may be what makes items redundant.
    items = remove_redundant_items(items)
    if len(items) > 3:
        return [AnyType()]
    else:
        return items


def simplify_recursive(typ):
    # type: (AbstractType) -> AbstractType
    """Simplify all components of a type."""
    if isinstance(typ, UnionType):
        return combine_types(typ.items)
    elif isinstance(typ, ClassType):
        simplified = ClassType(typ.name, [simplify_recursive(arg) for arg in typ.args])
        args = simplified.args
        if (simplified.name == 'Dict' and len(args) == 2
                and isinstance(args[0], ClassType) and args[0].name in ('str', 'Text')
                and isinstance(args[1], UnionType) and not is_optional(args[1])):
            # Looks like a potential case for TypedDict, which we don't properly support yet.
            return ClassType('Dict', [args[0], AnyType()])
        return simplified
    elif isinstance(typ, TupleType):
        return TupleType([simplify_recursive(item) for item in typ.items])
    return typ


def flatten_types(types):
    # type: (Iterable[AbstractType]) -> List[AbstractType]
    flattened = []
    for item in types:
        if not isinstance(item, UnionType):
            flattened.append(item)
        else:
            flattened.extend(flatten_types(item.items))
    return flattened


def dedupe_types(types):
    # type: (Iterable[AbstractType]) -> List[AbstractType]
    return sorted(set(types), key=lambda t: str(t))

def filter_ignored_items(items):
     # type: (List[AbstractType]) -> List[AbstractType]
    result = [item for item in items
              if not isinstance(item, ClassType) or
              item.name not in IGNORED_ITEMS]
    return result or [AnyType()]

def remove_redundant_items(items):
    # type: (List[AbstractType]) -> List[AbstractType]
    """Filter out redundant union items."""
    result = []
    for item in items:
        for other in items:
            if item is not other and is_redundant_union_item(item, other):
                break
        else:
            result.append(item)
    return result


def is_redundant_union_item(first, other):
    # type: (AbstractType, AbstractType) -> bool
    """If union has both items, is the first one redundant?

    For example, if first is 'str' and the other is 'Text', return True.

    If items are equal, return False.
    """
    if isinstance(first, ClassType) and isinstance(other, ClassType):
        if first.name == 'str' and other.name == 'Text':
            return True
        elif first.name == 'bool' and other.name == 'int':
            return True
        elif first.name == 'int' and other.name == 'float':
            return True
        elif (first.name in ('List', 'Dict', 'Set') and
                  other.name == first.name):
            if not first.args and other.args:
                return True
            elif len(first.args) == len(other.args) and first.args:
                result = all(first_arg == other_arg or other_arg == AnyType()
                             for first_arg, other_arg
                             in zip(first.args, other.args))
                return result

    return False


def merge_items(items):
    # type: (List[AbstractType]) -> List[AbstractType]
    """Merge union items that can be merged."""
    result = []
    while items:
        item = items.pop()
        merged = None
        for i, other in enumerate(items):
            merged = merged_type(item, other)
            if merged:
                break
        if merged:
            del items[i]
            items.append(merged)
        else:
            result.append(item)
    return list(reversed(result))


def merged_type(t, s):
    # type: (AbstractType, AbstractType) -> Optional[AbstractType]
    """Return merged type if two items can be merged in to a different, more general type.

    Return None if merging is not possible.
    """
    if isinstance(t, TupleType) and isinstance(s, TupleType):
        if len(t.items) == len(s.items):
            return TupleType([combine_types([ti, si]) for ti, si in zip(t.items, s.items)])
        all_items = t.items + s.items
        if all_items and all(item == all_items[0] for item in all_items[1:]):
            # Merge multiple compatible fixed-length tuples into a variable-length tuple type.
            return ClassType('Tuple', [all_items[0]])
    elif (isinstance(t, TupleType) and isinstance(s, ClassType) and s.name == 'Tuple'
          and len(s.args) == 1):
        if all(item == s.args[0] for item in t.items):
            # Merge fixed-length tuple and variable-length tuple.
            return s
    elif isinstance(s, TupleType) and isinstance(t, ClassType) and t.name == 'Tuple':
        return merged_type(s, t)
    elif isinstance(s, NoReturnType):
        return t
    elif isinstance(t, NoReturnType):
        return s
    elif isinstance(s, AnyType):
        # This seems to be usually desirable, since Anys tend to come from unknown types.
        return t
    elif isinstance(t, AnyType):
        # Similar to above.
        return s
    return None