0
# PyTD Type System
1

2
Python Type Declaration system for representing, manipulating, and optimizing type information. PyTD provides the foundational data structures and algorithms that enable PyType's type inference, stub generation, and type checking capabilities.
3

4
## Capabilities
5

6
### Core PyTD Types
7

8
Fundamental AST node types that represent Python type declarations and form the building blocks of PyType's type system.
9

10
```python { .api }
11
class Node:
12
    """Base class for all PyTD AST nodes."""
13
    
14
class Type(Node):
15
    """Base marker class for type representation nodes."""
16

17
class TypeDeclUnit(Node):
18
    """
19
    Top-level module node containing all declarations.
20
    
21
    Represents a complete Python module with its type information,
22
    including classes, functions, constants, and imports.
23
    """
24
    
25
    def __init__(self, name, constants, type_params, classes, functions, aliases, modules):
26
        """
27
        Initialize module type declaration.
28
        
29
        Parameters:
30
        - name (str): Module name
31
        - constants (list): Module-level constants
32
        - type_params (list): Type parameters  
33
        - classes (list): Class declarations
34
        - functions (list): Function declarations
35
        - aliases (list): Type aliases
36
        - modules (list): Imported modules
37
        """
38

39
class Constant(Node):
40
    """
41
    Module-level constant declaration.
42
    
43
    Represents constants defined at module scope with their inferred types.
44
    """
45
    
46
    def __init__(self, name, type):
47
        """
48
        Initialize constant declaration.
49
        
50
        Parameters:
51
        - name (str): Constant name
52
        - type (Type): Inferred type of constant
53
        """
54

55
class Alias(Node):
56
    """
57
    Type alias or symbolic link to classes in other modules.
58
    
59
    Enables references to types defined elsewhere without full qualification.
60
    """
61
    
62
    def __init__(self, name, type):
63
        """
64
        Initialize type alias.
65
        
66
        Parameters:
67
        - name (str): Alias name
68
        - type (Type): Target type being aliased
69
        """
70

71
class Class(Node):
72
    """
73
    Class declaration with methods, properties, and inheritance.
74
    
75
    Represents complete class type information including method signatures,
76
    class hierarchy, and special methods.
77
    """
78
    
79
    def __init__(self, name, metaclass, parents, methods, constants, classes, decorators, slots, template):
80
        """
81
        Initialize class declaration.
82
        
83
        Parameters:
84
        - name (str): Class name
85
        - metaclass (Type): Metaclass type
86
        - parents (list): Parent class types
87
        - methods (list): Method declarations
88
        - constants (list): Class constants
89
        - classes (list): Nested classes
90
        - decorators (list): Class decorators
91
        - slots (list): __slots__ specification
92
        - template (list): Generic type parameters
93
        """
94

95
class Function(Node):
96
    """
97
    Function declaration with signature and type information.
98
    
99
    Represents function signatures including parameters, return types,
100
    and special function properties.
101
    """
102
    
103
    def __init__(self, name, signatures, kind, flags):
104
        """
105
        Initialize function declaration.
106
        
107
        Parameters:
108
        - name (str): Function name
109
        - signatures (list): Function signatures
110
        - kind (str): Function kind (method, classmethod, staticmethod)
111
        - flags (int): Function property flags
112
        """
113
```
114

115
### Type Representation Classes
116

117
Specialized classes for representing different kinds of Python types in the PyTD system.
118

119
```python { .api }
120
class GenericType(Type):
121
    """Generic type with type parameters (e.g., List[int])."""
122
    
123
class ClassType(Type):
124
    """Reference to a class type."""
125
    
126
class UnionType(Type):
127
    """Union of multiple types (e.g., int | str)."""
128
    
129
class CallableType(Type):
130
    """Callable type with parameter and return types."""
131
    
132
class TupleType(Type):
133
    """Tuple type with element types."""
134
    
135
class LiteralType(Type):
136
    """Literal type for specific values."""
137
    
138
class TypeParameter(Type):
139
    """Type parameter for generic types."""
140
```
141

142
### PyTD Utilities
143

144
Essential utilities for working with PyTD AST structures, including printing, concatenation, and manipulation functions.
145

146
```python { .api }
147
def Print(ast, multiline_args=False):
148
    """
149
    Pretty print PyTD AST to string representation.
150
    
151
    Parameters:
152
    - ast (Node): PyTD AST node to print
153
    - multiline_args (bool): Whether to format arguments across multiple lines
154
    
155
    Returns:
156
    str: String representation of the AST
157
    """
158

159
def Concat(*asts):
160
    """
161
    Concatenate multiple PyTD ASTs into a single TypeDeclUnit.
162
    
163
    Parameters:
164
    - *asts: Variable number of TypeDeclUnit nodes to concatenate
165
    
166
    Returns:
167
    TypeDeclUnit: Combined AST containing all input declarations
168
    """
169
```
170

171
Example usage:
172

173
```python
174
from pytype.pytd import pytd_utils, pytd
175

176
# Create function signature
177
signature = pytd.Signature(
178
    params=[
179
        pytd.Parameter("x", pytd.NamedType("int")),
180
        pytd.Parameter("y", pytd.NamedType("int"))
181
    ],
182
    return_type=pytd.NamedType("int")
183
)
184

185
# Create function
186
func = pytd.Function(
187
    name="add",
188
    signatures=[signature],
189
    kind=pytd.MethodKind.METHOD,
190
    flags=0
191
)
192

193
# Print the function
194
print(pytd_utils.Print(func))
195
# Output: def add(x: int, y: int) -> int: ...
196
```
197

198
### PyTD Optimization
199

200
Optimizes PyTD ASTs for performance and correctness, including type simplification, dead code elimination, and union optimization.
201

202
```python { .api }
203
def Optimize(ast, builtins, lossy=False, use_abcs=False, max_union=4, remove_mutable=False, can_do_lookup=True):
204
    """
205
    Optimize PyTD AST for performance and correctness.
206
    
207
    Parameters:
208
    - ast (TypeDeclUnit): PyTD AST to optimize
209
    - builtins (TypeDeclUnit): Built-in types for reference
210
    - lossy (bool): Whether to allow lossy optimizations
211
    - use_abcs (bool): Whether to use abstract base classes
212
    - max_union (int): Maximum union size before simplification
213
    - remove_mutable (bool): Whether to remove mutable types
214
    - can_do_lookup (bool): Whether type lookups are available
215
    
216
    Returns:
217
    TypeDeclUnit: Optimized PyTD AST
218
    """
219
```
220

221
Example usage:
222

223
```python
224
from pytype.pytd import optimize
225
from pytype import io
226

227
# Generate AST from source
228
source = '''
229
def process_data(items):
230
    result = []
231
    for item in items:
232
        result.append(item.upper())
233
    return result
234
'''
235

236
ast = io.generate_pyi_ast(source)
237

238
# Optimize the AST
239
builtins_ast = # ... load builtins
240
optimized_ast = optimize.Optimize(
241
    ast, 
242
    builtins_ast,
243
    max_union=6,
244
    use_abcs=True
245
)
246
```
247

248
### PyTD Visitors
249

250
Visitor pattern implementation for traversing and transforming PyTD ASTs, enabling custom analysis and modification operations.
251

252
```python { .api }
253
class Visitor:
254
    """
255
    Base visitor class for PyTD AST traversal.
256
    
257
    Provides generic visitor pattern implementation with pre- and post-order
258
    traversal hooks for all PyTD node types.
259
    """
260
    
261
    def Visit(self, node):
262
        """
263
        Visit a PyTD node and its children.
264
        
265
        Parameters:
266
        - node (Node): PyTD node to visit
267
        
268
        Returns:
269
        Node: Potentially modified node
270
        """
271
    
272
    def VisitTypeDeclUnit(self, node):
273
        """Visit TypeDeclUnit node."""
274
        
275
    def VisitClass(self, node):
276
        """Visit Class node."""
277
        
278
    def VisitFunction(self, node):
279
        """Visit Function node."""
280
    
281
    # ... additional visit methods for each node type
282
```
283

284
Common visitor implementations include:
285

286
```python { .api }
287
class ReplaceTypes(Visitor):
288
    """Replace specific types throughout AST."""
289
    
290
class CollectDependencies(Visitor):
291
    """Collect type dependencies from AST."""
292
    
293
class VerifyVisitor(Visitor):
294
    """Verify AST correctness and consistency."""
295
```
296

297
### Module and Import Handling
298

299
Classes for representing module imports and cross-module type references.
300

301
```python { .api }
302
class Module(Node):
303
    """
304
    Imported module representation.
305
    
306
    Represents a module import with its qualified name and alias.
307
    """
308
    
309
    def __init__(self, module_name, alias=None):
310
        """
311
        Initialize module import.
312
        
313
        Parameters:
314
        - module_name (str): Fully qualified module name
315
        - alias (str, optional): Import alias
316
        """
317
```
318

319
### Working with PyTD Files
320

321
```python
322
from pytype.pytd import pytd, pytd_utils
323
from pytype.pyi import parser
324

325
# Parse .pyi file to PyTD AST
326
with open("types.pyi", "r") as f:
327
    pyi_content = f.read()
328

329
ast = parser.parse_string(pyi_content, filename="types.pyi")
330

331
# Manipulate the AST
332
# ... modify classes, functions, etc.
333

334
# Convert back to string
335
pyi_output = pytd_utils.Print(ast)
336

337
# Write optimized version
338
with open("optimized_types.pyi", "w") as f:
339
    f.write(pyi_output)
340
```
341

342
### Error Handling
343

344
PyTD operations can raise various exceptions during AST construction, optimization, and printing:
345

346
```python
347
from pytype.pytd import pytd
348

349
try:
350
    # Create PyTD structures
351
    ast = pytd.TypeDeclUnit(...)
352
    optimized = optimize.Optimize(ast, builtins)
353
    output = pytd_utils.Print(optimized)
354
except pytd.ParseError as e:
355
    print(f"PyTD parsing error: {e}")
356
except Exception as e:
357
    print(f"PyTD processing error: {e}")
358
```