brainstate 0.2.1__py2.py3-none-any.whl → 0.2.2__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (115) hide show
  1. brainstate/__init__.py +167 -169
  2. brainstate/_compatible_import.py +340 -340
  3. brainstate/_compatible_import_test.py +681 -681
  4. brainstate/_deprecation.py +210 -210
  5. brainstate/_deprecation_test.py +2297 -2319
  6. brainstate/_error.py +45 -45
  7. brainstate/_state.py +2157 -1652
  8. brainstate/_state_test.py +1129 -52
  9. brainstate/_utils.py +47 -47
  10. brainstate/environ.py +1495 -1495
  11. brainstate/environ_test.py +1223 -1223
  12. brainstate/graph/__init__.py +22 -22
  13. brainstate/graph/_node.py +240 -240
  14. brainstate/graph/_node_test.py +589 -589
  15. brainstate/graph/_operation.py +1620 -1624
  16. brainstate/graph/_operation_test.py +1147 -1147
  17. brainstate/mixin.py +1447 -1433
  18. brainstate/mixin_test.py +1017 -1017
  19. brainstate/nn/__init__.py +146 -137
  20. brainstate/nn/_activations.py +1100 -1100
  21. brainstate/nn/_activations_test.py +354 -354
  22. brainstate/nn/_collective_ops.py +635 -633
  23. brainstate/nn/_collective_ops_test.py +774 -774
  24. brainstate/nn/_common.py +226 -226
  25. brainstate/nn/_common_test.py +134 -154
  26. brainstate/nn/_conv.py +2010 -2010
  27. brainstate/nn/_conv_test.py +849 -849
  28. brainstate/nn/_delay.py +575 -575
  29. brainstate/nn/_delay_test.py +243 -243
  30. brainstate/nn/_dropout.py +618 -618
  31. brainstate/nn/_dropout_test.py +480 -477
  32. brainstate/nn/_dynamics.py +870 -1267
  33. brainstate/nn/_dynamics_test.py +53 -67
  34. brainstate/nn/_elementwise.py +1298 -1298
  35. brainstate/nn/_elementwise_test.py +829 -829
  36. brainstate/nn/_embedding.py +408 -408
  37. brainstate/nn/_embedding_test.py +156 -156
  38. brainstate/nn/_event_fixedprob.py +233 -233
  39. brainstate/nn/_event_fixedprob_test.py +115 -115
  40. brainstate/nn/_event_linear.py +83 -83
  41. brainstate/nn/_event_linear_test.py +121 -121
  42. brainstate/nn/_exp_euler.py +254 -254
  43. brainstate/nn/_exp_euler_test.py +377 -377
  44. brainstate/nn/_linear.py +744 -744
  45. brainstate/nn/_linear_test.py +475 -475
  46. brainstate/nn/_metrics.py +1070 -1070
  47. brainstate/nn/_metrics_test.py +611 -611
  48. brainstate/nn/_module.py +391 -384
  49. brainstate/nn/_module_test.py +40 -40
  50. brainstate/nn/_normalizations.py +1334 -1334
  51. brainstate/nn/_normalizations_test.py +699 -699
  52. brainstate/nn/_paddings.py +1020 -1020
  53. brainstate/nn/_paddings_test.py +722 -722
  54. brainstate/nn/_poolings.py +2239 -2239
  55. brainstate/nn/_poolings_test.py +952 -952
  56. brainstate/nn/_rnns.py +946 -946
  57. brainstate/nn/_rnns_test.py +592 -592
  58. brainstate/nn/_utils.py +216 -216
  59. brainstate/nn/_utils_test.py +401 -401
  60. brainstate/nn/init.py +809 -809
  61. brainstate/nn/init_test.py +180 -180
  62. brainstate/random/__init__.py +270 -270
  63. brainstate/random/{_rand_funs.py → _fun.py} +3938 -3938
  64. brainstate/random/{_rand_funs_test.py → _fun_test.py} +638 -640
  65. brainstate/random/_impl.py +672 -0
  66. brainstate/random/{_rand_seed.py → _seed.py} +675 -675
  67. brainstate/random/{_rand_seed_test.py → _seed_test.py} +48 -48
  68. brainstate/random/{_rand_state.py → _state.py} +1320 -1617
  69. brainstate/random/{_rand_state_test.py → _state_test.py} +551 -551
  70. brainstate/transform/__init__.py +56 -59
  71. brainstate/transform/_ad_checkpoint.py +176 -176
  72. brainstate/transform/_ad_checkpoint_test.py +49 -49
  73. brainstate/transform/_autograd.py +1025 -1025
  74. brainstate/transform/_autograd_test.py +1289 -1289
  75. brainstate/transform/_conditions.py +316 -316
  76. brainstate/transform/_conditions_test.py +220 -220
  77. brainstate/transform/_error_if.py +94 -94
  78. brainstate/transform/_error_if_test.py +52 -52
  79. brainstate/transform/_find_state.py +200 -0
  80. brainstate/transform/_find_state_test.py +84 -0
  81. brainstate/transform/_jit.py +399 -399
  82. brainstate/transform/_jit_test.py +143 -143
  83. brainstate/transform/_loop_collect_return.py +675 -675
  84. brainstate/transform/_loop_collect_return_test.py +58 -58
  85. brainstate/transform/_loop_no_collection.py +283 -283
  86. brainstate/transform/_loop_no_collection_test.py +50 -50
  87. brainstate/transform/_make_jaxpr.py +2176 -2016
  88. brainstate/transform/_make_jaxpr_test.py +1634 -1510
  89. brainstate/transform/_mapping.py +607 -529
  90. brainstate/transform/_mapping_test.py +104 -194
  91. brainstate/transform/_progress_bar.py +255 -255
  92. brainstate/transform/_unvmap.py +256 -256
  93. brainstate/transform/_util.py +286 -286
  94. brainstate/typing.py +837 -837
  95. brainstate/typing_test.py +780 -780
  96. brainstate/util/__init__.py +27 -27
  97. brainstate/util/_others.py +1024 -1024
  98. brainstate/util/_others_test.py +962 -962
  99. brainstate/util/_pretty_pytree.py +1301 -1301
  100. brainstate/util/_pretty_pytree_test.py +675 -675
  101. brainstate/util/_pretty_repr.py +462 -462
  102. brainstate/util/_pretty_repr_test.py +696 -696
  103. brainstate/util/filter.py +945 -945
  104. brainstate/util/filter_test.py +911 -911
  105. brainstate/util/struct.py +910 -910
  106. brainstate/util/struct_test.py +602 -602
  107. {brainstate-0.2.1.dist-info → brainstate-0.2.2.dist-info}/METADATA +108 -108
  108. brainstate-0.2.2.dist-info/RECORD +111 -0
  109. {brainstate-0.2.1.dist-info → brainstate-0.2.2.dist-info}/licenses/LICENSE +202 -202
  110. brainstate/transform/_eval_shape.py +0 -145
  111. brainstate/transform/_eval_shape_test.py +0 -38
  112. brainstate/transform/_random.py +0 -171
  113. brainstate-0.2.1.dist-info/RECORD +0 -111
  114. {brainstate-0.2.1.dist-info → brainstate-0.2.2.dist-info}/WHEEL +0 -0
  115. {brainstate-0.2.1.dist-info → brainstate-0.2.2.dist-info}/top_level.txt +0 -0
@@ -1,462 +1,462 @@
1
- # The file is adapted from the Flax library (https://github.com/google/flax).
2
- # The credit should go to the Flax authors.
3
- #
4
- # Copyright 2024 The Flax Authors.
5
- #
6
- # Licensed under the Apache License, Version 2.0 (the "License");
7
- # you may not use this file except in compliance with the License.
8
- # You may obtain a copy of the License at
9
- #
10
- # http://www.apache.org/licenses/LICENSE-2.0
11
- #
12
- # Unless required by applicable law or agreed to in writing, software
13
- # distributed under the License is distributed on an "AS IS" BASIS,
14
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
- # See the License for the specific language governing permissions and
16
- # limitations under the License.
17
-
18
- """
19
- Pretty representation utilities for creating human-readable string representations.
20
-
21
- This module provides utilities for creating customizable pretty representations of
22
- objects, with support for nested structures and circular reference detection.
23
- """
24
-
25
- import dataclasses
26
- import threading
27
- from abc import ABC, abstractmethod
28
- from functools import partial
29
- from typing import Any, Iterator, Mapping, TypeVar, Union, Callable, Optional
30
-
31
- __all__ = [
32
- 'yield_unique_pretty_repr_items',
33
- 'PrettyType',
34
- 'PrettyAttr',
35
- 'PrettyRepr',
36
- 'PrettyMapping',
37
- 'MappingReprMixin',
38
- ]
39
-
40
- A = TypeVar('A')
41
- B = TypeVar('B')
42
-
43
-
44
- @dataclasses.dataclass
45
- class PrettyType:
46
- """
47
- Configuration for pretty representation of objects.
48
-
49
- Attributes
50
- ----------
51
- type : Union[str, type]
52
- The type name or type object to display.
53
- start : str, default='('
54
- The opening delimiter for the representation.
55
- end : str, default=')'
56
- The closing delimiter for the representation.
57
- value_sep : str, default='='
58
- The separator between keys and values.
59
- elem_indent : str, default=' '
60
- The indentation for nested elements.
61
- empty_repr : str, default=''
62
- The representation for empty objects.
63
- """
64
- type: Union[str, type]
65
- start: str = '('
66
- end: str = ')'
67
- value_sep: str = '='
68
- elem_indent: str = ' '
69
- empty_repr: str = ''
70
-
71
-
72
- @dataclasses.dataclass
73
- class PrettyAttr:
74
- """
75
- Configuration for pretty representation of attributes.
76
-
77
- Attributes
78
- ----------
79
- key : str
80
- The attribute name or key.
81
- value : Union[str, Any]
82
- The attribute value.
83
- start : str, default=''
84
- Optional prefix for the attribute.
85
- end : str, default=''
86
- Optional suffix for the attribute.
87
- """
88
- key: str
89
- value: Union[str, Any]
90
- start: str = ''
91
- end: str = ''
92
-
93
-
94
- class PrettyRepr(ABC):
95
- """
96
- Interface for pretty representation of objects.
97
-
98
- This abstract base class provides a framework for creating custom
99
- pretty representations of objects by yielding PrettyType and PrettyAttr
100
- instances.
101
-
102
- Examples
103
- --------
104
- .. code-block:: python
105
-
106
- >>> class MyObject(PrettyRepr):
107
- ... def __init__(self, key, value):
108
- ... self.key = key
109
- ... self.value = value
110
- ...
111
- ... def __pretty_repr__(self):
112
- ... yield PrettyType(type='MyObject', start='{', end='}')
113
- ... yield PrettyAttr('key', self.key)
114
- ... yield PrettyAttr('value', self.value)
115
- ...
116
- >>> obj = MyObject('foo', 42)
117
- >>> print(obj)
118
- MyObject{
119
- key=foo,
120
- value=42
121
- }
122
- """
123
- __slots__ = ()
124
-
125
- @abstractmethod
126
- def __pretty_repr__(self) -> Iterator[Union[PrettyType, PrettyAttr]]:
127
- """
128
- Generate the pretty representation of the object.
129
-
130
- Yields
131
- ------
132
- Union[PrettyType, PrettyAttr]
133
- First yield should be PrettyType, followed by PrettyAttr instances.
134
- """
135
- raise NotImplementedError
136
-
137
- def __repr__(self) -> str:
138
- """
139
- Generate string representation using pretty representation.
140
-
141
- Returns
142
- -------
143
- str
144
- The formatted string representation of the object.
145
- """
146
- return pretty_repr_object(self)
147
-
148
-
149
- def pretty_repr_elem(obj: PrettyType, elem: Any) -> str:
150
- """
151
- Constructs a string representation of a single element within a pretty representation.
152
-
153
- This function takes a `PrettyType` object and an element, which must be an instance
154
- of `PrettyAttr`, and generates a formatted string that represents the element. The
155
- formatting is based on the configuration provided by the `PrettyType` object.
156
-
157
- Parameters
158
- ----------
159
- obj : PrettyType
160
- The configuration object that defines how the element should be formatted.
161
- It includes details such as indentation, separators, and surrounding characters.
162
- elem : Any
163
- The element to be represented. It must be an instance of `PrettyAttr`, which
164
- contains the key and value to be formatted.
165
-
166
- Returns
167
- -------
168
- str
169
- A string that represents the element in a formatted manner, adhering to the
170
- configuration specified by the `PrettyType` object.
171
-
172
- Raises
173
- ------
174
- TypeError
175
- If the provided element is not an instance of `PrettyAttr`.
176
- """
177
- if not isinstance(elem, PrettyAttr):
178
- raise TypeError(f'Item must be Elem, got {type(elem).__name__}')
179
-
180
- value = elem.value if isinstance(elem.value, str) else repr(elem.value)
181
- value = value.replace('\n', '\n' + obj.elem_indent)
182
-
183
- return f'{obj.elem_indent}{elem.start}{elem.key}{obj.value_sep}{value}{elem.end}'
184
-
185
-
186
- def pretty_repr_object(obj: PrettyRepr) -> str:
187
- """
188
- Generates a pretty string representation of an object that implements the PrettyRepr interface.
189
-
190
- This function utilizes the __pretty_repr__ method of the PrettyRepr interface to obtain
191
- a structured representation of the object, which includes both the type and attributes
192
- of the object in a human-readable format.
193
-
194
- Parameters
195
- ----------
196
- obj : PrettyRepr
197
- The object for which the pretty representation is to be generated. The object must
198
- implement the PrettyRepr interface.
199
-
200
- Returns
201
- -------
202
- str
203
- A string that represents the object in a pretty format, including its type and attributes.
204
- The format is determined by the PrettyType and PrettyAttr instances yielded by the
205
- __pretty_repr__ method of the object.
206
-
207
- Raises
208
- ------
209
- TypeError
210
- If the provided object does not implement the PrettyRepr interface or if the first item
211
- yielded by the __pretty_repr__ method is not an instance of PrettyType.
212
- """
213
- if not isinstance(obj, PrettyRepr):
214
- raise TypeError(f'Object {obj!r} is not representable')
215
-
216
- iterator = obj.__pretty_repr__()
217
- obj_repr = next(iterator)
218
-
219
- # repr object
220
- if not isinstance(obj_repr, PrettyType):
221
- raise TypeError(f'First item must be PrettyType, got {type(obj_repr).__name__}')
222
-
223
- # repr attributes
224
- elem_reprs = tuple(map(partial(pretty_repr_elem, obj_repr), iterator))
225
- elems = ',\n'.join(elem_reprs)
226
- if elems:
227
- elems = '\n' + elems + '\n'
228
- else:
229
- elems = obj_repr.empty_repr
230
-
231
- # repr object type
232
- type_repr = obj_repr.type if isinstance(obj_repr.type, str) else obj_repr.type.__name__
233
-
234
- # return repr
235
- return f'{type_repr}{obj_repr.start}{elems}{obj_repr.end}'
236
-
237
-
238
- class MappingReprMixin(Mapping[A, B]):
239
- """
240
- Mapping mixin for pretty representation.
241
-
242
- This mixin provides a default pretty representation for mapping-like objects.
243
-
244
- Examples
245
- --------
246
- .. code-block:: python
247
-
248
- >>> class MyMapping(dict, MappingReprMixin):
249
- ... pass
250
- ...
251
- >>> m = MyMapping({'a': 1, 'b': 2})
252
- >>> print(m)
253
- {
254
- 'a': 1,
255
- 'b': 2
256
- }
257
- """
258
-
259
- def __pretty_repr__(self) -> Iterator[Union[PrettyType, PrettyAttr]]:
260
- """
261
- Generate pretty representation for mapping.
262
-
263
- Yields
264
- ------
265
- Union[PrettyType, PrettyAttr]
266
- PrettyType followed by PrettyAttr for each key-value pair.
267
- """
268
- yield PrettyType(type='', value_sep=': ', start='{', end='}')
269
-
270
- for key, value in self.items():
271
- yield PrettyAttr(repr(key), value)
272
-
273
-
274
- @dataclasses.dataclass(repr=False)
275
- class PrettyMapping(PrettyRepr):
276
- """
277
- Pretty representation of a mapping.
278
-
279
- Attributes
280
- ----------
281
- mapping : Mapping
282
- The mapping to represent.
283
- type_name : str, default=''
284
- Optional type name to display.
285
-
286
- Examples
287
- --------
288
- .. code-block:: python
289
-
290
- >>> m = PrettyMapping({'a': 1, 'b': 2}, type_name='MyDict')
291
- >>> print(m)
292
- MyDict{
293
- 'a': 1,
294
- 'b': 2
295
- }
296
- """
297
- mapping: Mapping
298
- type_name: str = ''
299
-
300
- def __pretty_repr__(self) -> Iterator[Union[PrettyType, PrettyAttr]]:
301
- """
302
- Generate pretty representation for the mapping.
303
-
304
- Yields
305
- ------
306
- Union[PrettyType, PrettyAttr]
307
- PrettyType followed by PrettyAttr for each key-value pair.
308
- """
309
- yield PrettyType(type=self.type_name, value_sep=': ', start='{', end='}')
310
-
311
- for key, value in self.mapping.items():
312
- yield PrettyAttr(repr(key), value)
313
-
314
-
315
- @dataclasses.dataclass
316
- class PrettyReprContext(threading.local):
317
- """
318
- A thread-local context for managing the state of pretty representation.
319
-
320
- This class is used to keep track of objects that have been seen during
321
- the generation of pretty representations, preventing infinite recursion
322
- in cases of circular references.
323
-
324
- Attributes
325
- ----------
326
- seen_modules_repr : dict[int, Any] | None
327
- A dictionary mapping object IDs to objects that have been seen
328
- during the pretty representation process. This is used to avoid
329
- representing the same object multiple times.
330
- """
331
- seen_modules_repr: dict[int, Any] | None = None
332
-
333
-
334
- CONTEXT = PrettyReprContext()
335
-
336
-
337
- def _default_repr_object(node: Any) -> Iterator[PrettyType]:
338
- """
339
- Generate a default pretty representation for an object.
340
-
341
- This function yields a `PrettyType` instance that represents the type
342
- of the given object. It is used as a default method for representing
343
- objects when no custom representation function is provided.
344
-
345
- Parameters
346
- ----------
347
- node : Any
348
- The object for which the pretty representation is to be generated.
349
-
350
- Yields
351
- ------
352
- PrettyType
353
- An instance of `PrettyType` that contains the type information of
354
- the object.
355
- """
356
- yield PrettyType(type=type(node))
357
-
358
-
359
- def _default_repr_attr(node: Any) -> Iterator[PrettyAttr]:
360
- """
361
- Generate a default pretty representation for the attributes of an object.
362
-
363
- This function iterates over the attributes of the given object and yields
364
- a `PrettyAttr` instance for each attribute that does not start with an
365
- underscore. The `PrettyAttr` instances contain the attribute name and its
366
- string representation.
367
-
368
- Parameters
369
- ----------
370
- node : Any
371
- The object whose attributes are to be represented.
372
-
373
- Yields
374
- ------
375
- PrettyAttr
376
- An instance of `PrettyAttr` for each non-private attribute of the object,
377
- containing the attribute name and its string representation.
378
- """
379
- for name, value in vars(node).items():
380
- if name.startswith('_'):
381
- continue
382
- yield PrettyAttr(name, repr(value))
383
-
384
-
385
- def yield_unique_pretty_repr_items(
386
- node: Any,
387
- repr_object: Optional[Callable] = None,
388
- repr_attr: Optional[Callable] = None
389
- ) -> Iterator[Union[PrettyType, PrettyAttr]]:
390
- """
391
- Generate a pretty representation of an object while avoiding duplicate representations.
392
-
393
- This function yields a structured representation of an object, using custom or default
394
- methods for representing the object itself and its attributes. It ensures that each
395
- object is only represented once to prevent infinite recursion in cases of circular
396
- references.
397
-
398
- Parameters
399
- ----------
400
- node : Any
401
- The object to be represented.
402
- repr_object : Optional[Callable], optional
403
- A callable that yields the representation of the object itself.
404
- If not provided, a default representation function is used.
405
- repr_attr : Optional[Callable], optional
406
- A callable that yields the representation of the object's attributes.
407
- If not provided, a default attribute representation function is used.
408
-
409
- Yields
410
- ------
411
- Union[PrettyType, PrettyAttr]
412
- The pretty representation of the object and its attributes,
413
- avoiding duplicates by tracking seen objects.
414
-
415
- Examples
416
- --------
417
- .. code-block:: python
418
-
419
- >>> class Node:
420
- ... def __init__(self, value, next=None):
421
- ... self.value = value
422
- ... self.next = next
423
- ...
424
- >>> # Create circular reference
425
- >>> node1 = Node(1)
426
- >>> node2 = Node(2, node1)
427
- >>> node1.next = node2
428
- ...
429
- >>> # This will handle circular reference gracefully
430
- >>> for item in yield_unique_pretty_repr_items(node1):
431
- ... print(item)
432
- """
433
- if repr_object is None:
434
- repr_object = _default_repr_object
435
- if repr_attr is None:
436
- repr_attr = _default_repr_attr
437
-
438
- if CONTEXT.seen_modules_repr is None:
439
- # CONTEXT.seen_modules_repr = set()
440
- CONTEXT.seen_modules_repr = dict()
441
- clear_seen = True
442
- else:
443
- clear_seen = False
444
-
445
- # Avoid infinite recursion
446
- if id(node) in CONTEXT.seen_modules_repr:
447
- yield PrettyType(type=type(node), empty_repr='...')
448
- return
449
-
450
- # repr object
451
- yield from repr_object(node)
452
-
453
- # Add to seen modules
454
- # CONTEXT.seen_modules_repr.add(id(node))
455
- CONTEXT.seen_modules_repr[id(node)] = node
456
-
457
- try:
458
- # repr attributes
459
- yield from repr_attr(node)
460
- finally:
461
- if clear_seen:
462
- CONTEXT.seen_modules_repr = None
1
+ # The file is adapted from the Flax library (https://github.com/google/flax).
2
+ # The credit should go to the Flax authors.
3
+ #
4
+ # Copyright 2024 The Flax Authors.
5
+ #
6
+ # Licensed under the Apache License, Version 2.0 (the "License");
7
+ # you may not use this file except in compliance with the License.
8
+ # You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing, software
13
+ # distributed under the License is distributed on an "AS IS" BASIS,
14
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
+ # See the License for the specific language governing permissions and
16
+ # limitations under the License.
17
+
18
+ """
19
+ Pretty representation utilities for creating human-readable string representations.
20
+
21
+ This module provides utilities for creating customizable pretty representations of
22
+ objects, with support for nested structures and circular reference detection.
23
+ """
24
+
25
+ import dataclasses
26
+ import threading
27
+ from abc import ABC, abstractmethod
28
+ from functools import partial
29
+ from typing import Any, Iterator, Mapping, TypeVar, Union, Callable, Optional
30
+
31
+ __all__ = [
32
+ 'yield_unique_pretty_repr_items',
33
+ 'PrettyType',
34
+ 'PrettyAttr',
35
+ 'PrettyRepr',
36
+ 'PrettyMapping',
37
+ 'MappingReprMixin',
38
+ ]
39
+
40
+ A = TypeVar('A')
41
+ B = TypeVar('B')
42
+
43
+
44
+ @dataclasses.dataclass
45
+ class PrettyType:
46
+ """
47
+ Configuration for pretty representation of objects.
48
+
49
+ Attributes
50
+ ----------
51
+ type : Union[str, type]
52
+ The type name or type object to display.
53
+ start : str, default='('
54
+ The opening delimiter for the representation.
55
+ end : str, default=')'
56
+ The closing delimiter for the representation.
57
+ value_sep : str, default='='
58
+ The separator between keys and values.
59
+ elem_indent : str, default=' '
60
+ The indentation for nested elements.
61
+ empty_repr : str, default=''
62
+ The representation for empty objects.
63
+ """
64
+ type: Union[str, type]
65
+ start: str = '('
66
+ end: str = ')'
67
+ value_sep: str = '='
68
+ elem_indent: str = ' '
69
+ empty_repr: str = ''
70
+
71
+
72
+ @dataclasses.dataclass
73
+ class PrettyAttr:
74
+ """
75
+ Configuration for pretty representation of attributes.
76
+
77
+ Attributes
78
+ ----------
79
+ key : str
80
+ The attribute name or key.
81
+ value : Union[str, Any]
82
+ The attribute value.
83
+ start : str, default=''
84
+ Optional prefix for the attribute.
85
+ end : str, default=''
86
+ Optional suffix for the attribute.
87
+ """
88
+ key: str
89
+ value: Union[str, Any]
90
+ start: str = ''
91
+ end: str = ''
92
+
93
+
94
+ class PrettyRepr(ABC):
95
+ """
96
+ Interface for pretty representation of objects.
97
+
98
+ This abstract base class provides a framework for creating custom
99
+ pretty representations of objects by yielding PrettyType and PrettyAttr
100
+ instances.
101
+
102
+ Examples
103
+ --------
104
+ .. code-block:: python
105
+
106
+ >>> class MyObject(PrettyRepr):
107
+ ... def __init__(self, key, value):
108
+ ... self.key = key
109
+ ... self.value = value
110
+ ...
111
+ ... def __pretty_repr__(self):
112
+ ... yield PrettyType(type='MyObject', start='{', end='}')
113
+ ... yield PrettyAttr('key', self.key)
114
+ ... yield PrettyAttr('value', self.value)
115
+ ...
116
+ >>> obj = MyObject('foo', 42)
117
+ >>> print(obj)
118
+ MyObject{
119
+ key=foo,
120
+ value=42
121
+ }
122
+ """
123
+ __slots__ = ()
124
+
125
+ @abstractmethod
126
+ def __pretty_repr__(self) -> Iterator[Union[PrettyType, PrettyAttr]]:
127
+ """
128
+ Generate the pretty representation of the object.
129
+
130
+ Yields
131
+ ------
132
+ Union[PrettyType, PrettyAttr]
133
+ First yield should be PrettyType, followed by PrettyAttr instances.
134
+ """
135
+ raise NotImplementedError
136
+
137
+ def __repr__(self) -> str:
138
+ """
139
+ Generate string representation using pretty representation.
140
+
141
+ Returns
142
+ -------
143
+ str
144
+ The formatted string representation of the object.
145
+ """
146
+ return pretty_repr_object(self)
147
+
148
+
149
+ def pretty_repr_elem(obj: PrettyType, elem: Any) -> str:
150
+ """
151
+ Constructs a string representation of a single element within a pretty representation.
152
+
153
+ This function takes a `PrettyType` object and an element, which must be an instance
154
+ of `PrettyAttr`, and generates a formatted string that represents the element. The
155
+ formatting is based on the configuration provided by the `PrettyType` object.
156
+
157
+ Parameters
158
+ ----------
159
+ obj : PrettyType
160
+ The configuration object that defines how the element should be formatted.
161
+ It includes details such as indentation, separators, and surrounding characters.
162
+ elem : Any
163
+ The element to be represented. It must be an instance of `PrettyAttr`, which
164
+ contains the key and value to be formatted.
165
+
166
+ Returns
167
+ -------
168
+ str
169
+ A string that represents the element in a formatted manner, adhering to the
170
+ configuration specified by the `PrettyType` object.
171
+
172
+ Raises
173
+ ------
174
+ TypeError
175
+ If the provided element is not an instance of `PrettyAttr`.
176
+ """
177
+ if not isinstance(elem, PrettyAttr):
178
+ raise TypeError(f'Item must be Elem, got {type(elem).__name__}')
179
+
180
+ value = elem.value if isinstance(elem.value, str) else repr(elem.value)
181
+ value = value.replace('\n', '\n' + obj.elem_indent)
182
+
183
+ return f'{obj.elem_indent}{elem.start}{elem.key}{obj.value_sep}{value}{elem.end}'
184
+
185
+
186
+ def pretty_repr_object(obj: PrettyRepr) -> str:
187
+ """
188
+ Generates a pretty string representation of an object that implements the PrettyRepr interface.
189
+
190
+ This function utilizes the __pretty_repr__ method of the PrettyRepr interface to obtain
191
+ a structured representation of the object, which includes both the type and attributes
192
+ of the object in a human-readable format.
193
+
194
+ Parameters
195
+ ----------
196
+ obj : PrettyRepr
197
+ The object for which the pretty representation is to be generated. The object must
198
+ implement the PrettyRepr interface.
199
+
200
+ Returns
201
+ -------
202
+ str
203
+ A string that represents the object in a pretty format, including its type and attributes.
204
+ The format is determined by the PrettyType and PrettyAttr instances yielded by the
205
+ __pretty_repr__ method of the object.
206
+
207
+ Raises
208
+ ------
209
+ TypeError
210
+ If the provided object does not implement the PrettyRepr interface or if the first item
211
+ yielded by the __pretty_repr__ method is not an instance of PrettyType.
212
+ """
213
+ if not isinstance(obj, PrettyRepr):
214
+ raise TypeError(f'Object {obj!r} is not representable')
215
+
216
+ iterator = obj.__pretty_repr__()
217
+ obj_repr = next(iterator)
218
+
219
+ # repr object
220
+ if not isinstance(obj_repr, PrettyType):
221
+ raise TypeError(f'First item must be PrettyType, got {type(obj_repr).__name__}')
222
+
223
+ # repr attributes
224
+ elem_reprs = tuple(map(partial(pretty_repr_elem, obj_repr), iterator))
225
+ elems = ',\n'.join(elem_reprs)
226
+ if elems:
227
+ elems = '\n' + elems + '\n'
228
+ else:
229
+ elems = obj_repr.empty_repr
230
+
231
+ # repr object type
232
+ type_repr = obj_repr.type if isinstance(obj_repr.type, str) else obj_repr.type.__name__
233
+
234
+ # return repr
235
+ return f'{type_repr}{obj_repr.start}{elems}{obj_repr.end}'
236
+
237
+
238
+ class MappingReprMixin(Mapping[A, B]):
239
+ """
240
+ Mapping mixin for pretty representation.
241
+
242
+ This mixin provides a default pretty representation for mapping-like objects.
243
+
244
+ Examples
245
+ --------
246
+ .. code-block:: python
247
+
248
+ >>> class MyMapping(dict, MappingReprMixin):
249
+ ... pass
250
+ ...
251
+ >>> m = MyMapping({'a': 1, 'b': 2})
252
+ >>> print(m)
253
+ {
254
+ 'a': 1,
255
+ 'b': 2
256
+ }
257
+ """
258
+
259
+ def __pretty_repr__(self) -> Iterator[Union[PrettyType, PrettyAttr]]:
260
+ """
261
+ Generate pretty representation for mapping.
262
+
263
+ Yields
264
+ ------
265
+ Union[PrettyType, PrettyAttr]
266
+ PrettyType followed by PrettyAttr for each key-value pair.
267
+ """
268
+ yield PrettyType(type='', value_sep=': ', start='{', end='}')
269
+
270
+ for key, value in self.items():
271
+ yield PrettyAttr(repr(key), value)
272
+
273
+
274
+ @dataclasses.dataclass(repr=False)
275
+ class PrettyMapping(PrettyRepr):
276
+ """
277
+ Pretty representation of a mapping.
278
+
279
+ Attributes
280
+ ----------
281
+ mapping : Mapping
282
+ The mapping to represent.
283
+ type_name : str, default=''
284
+ Optional type name to display.
285
+
286
+ Examples
287
+ --------
288
+ .. code-block:: python
289
+
290
+ >>> m = PrettyMapping({'a': 1, 'b': 2}, type_name='MyDict')
291
+ >>> print(m)
292
+ MyDict{
293
+ 'a': 1,
294
+ 'b': 2
295
+ }
296
+ """
297
+ mapping: Mapping
298
+ type_name: str = ''
299
+
300
+ def __pretty_repr__(self) -> Iterator[Union[PrettyType, PrettyAttr]]:
301
+ """
302
+ Generate pretty representation for the mapping.
303
+
304
+ Yields
305
+ ------
306
+ Union[PrettyType, PrettyAttr]
307
+ PrettyType followed by PrettyAttr for each key-value pair.
308
+ """
309
+ yield PrettyType(type=self.type_name, value_sep=': ', start='{', end='}')
310
+
311
+ for key, value in self.mapping.items():
312
+ yield PrettyAttr(repr(key), value)
313
+
314
+
315
+ @dataclasses.dataclass
316
+ class PrettyReprContext(threading.local):
317
+ """
318
+ A thread-local context for managing the state of pretty representation.
319
+
320
+ This class is used to keep track of objects that have been seen during
321
+ the generation of pretty representations, preventing infinite recursion
322
+ in cases of circular references.
323
+
324
+ Attributes
325
+ ----------
326
+ seen_modules_repr : dict[int, Any] | None
327
+ A dictionary mapping object IDs to objects that have been seen
328
+ during the pretty representation process. This is used to avoid
329
+ representing the same object multiple times.
330
+ """
331
+ seen_modules_repr: dict[int, Any] | None = None
332
+
333
+
334
+ CONTEXT = PrettyReprContext()
335
+
336
+
337
+ def _default_repr_object(node: Any) -> Iterator[PrettyType]:
338
+ """
339
+ Generate a default pretty representation for an object.
340
+
341
+ This function yields a `PrettyType` instance that represents the type
342
+ of the given object. It is used as a default method for representing
343
+ objects when no custom representation function is provided.
344
+
345
+ Parameters
346
+ ----------
347
+ node : Any
348
+ The object for which the pretty representation is to be generated.
349
+
350
+ Yields
351
+ ------
352
+ PrettyType
353
+ An instance of `PrettyType` that contains the type information of
354
+ the object.
355
+ """
356
+ yield PrettyType(type=type(node))
357
+
358
+
359
+ def _default_repr_attr(node: Any) -> Iterator[PrettyAttr]:
360
+ """
361
+ Generate a default pretty representation for the attributes of an object.
362
+
363
+ This function iterates over the attributes of the given object and yields
364
+ a `PrettyAttr` instance for each attribute that does not start with an
365
+ underscore. The `PrettyAttr` instances contain the attribute name and its
366
+ string representation.
367
+
368
+ Parameters
369
+ ----------
370
+ node : Any
371
+ The object whose attributes are to be represented.
372
+
373
+ Yields
374
+ ------
375
+ PrettyAttr
376
+ An instance of `PrettyAttr` for each non-private attribute of the object,
377
+ containing the attribute name and its string representation.
378
+ """
379
+ for name, value in vars(node).items():
380
+ if name.startswith('_'):
381
+ continue
382
+ yield PrettyAttr(name, repr(value))
383
+
384
+
385
+ def yield_unique_pretty_repr_items(
386
+ node: Any,
387
+ repr_object: Optional[Callable] = None,
388
+ repr_attr: Optional[Callable] = None
389
+ ) -> Iterator[Union[PrettyType, PrettyAttr]]:
390
+ """
391
+ Generate a pretty representation of an object while avoiding duplicate representations.
392
+
393
+ This function yields a structured representation of an object, using custom or default
394
+ methods for representing the object itself and its attributes. It ensures that each
395
+ object is only represented once to prevent infinite recursion in cases of circular
396
+ references.
397
+
398
+ Parameters
399
+ ----------
400
+ node : Any
401
+ The object to be represented.
402
+ repr_object : Optional[Callable], optional
403
+ A callable that yields the representation of the object itself.
404
+ If not provided, a default representation function is used.
405
+ repr_attr : Optional[Callable], optional
406
+ A callable that yields the representation of the object's attributes.
407
+ If not provided, a default attribute representation function is used.
408
+
409
+ Yields
410
+ ------
411
+ Union[PrettyType, PrettyAttr]
412
+ The pretty representation of the object and its attributes,
413
+ avoiding duplicates by tracking seen objects.
414
+
415
+ Examples
416
+ --------
417
+ .. code-block:: python
418
+
419
+ >>> class Node:
420
+ ... def __init__(self, value, next=None):
421
+ ... self.value = value
422
+ ... self.next = next
423
+ ...
424
+ >>> # Create circular reference
425
+ >>> node1 = Node(1)
426
+ >>> node2 = Node(2, node1)
427
+ >>> node1.next = node2
428
+ ...
429
+ >>> # This will handle circular reference gracefully
430
+ >>> for item in yield_unique_pretty_repr_items(node1):
431
+ ... print(item)
432
+ """
433
+ if repr_object is None:
434
+ repr_object = _default_repr_object
435
+ if repr_attr is None:
436
+ repr_attr = _default_repr_attr
437
+
438
+ if CONTEXT.seen_modules_repr is None:
439
+ # CONTEXT.seen_modules_repr = set()
440
+ CONTEXT.seen_modules_repr = dict()
441
+ clear_seen = True
442
+ else:
443
+ clear_seen = False
444
+
445
+ # Avoid infinite recursion
446
+ if id(node) in CONTEXT.seen_modules_repr:
447
+ yield PrettyType(type=type(node), empty_repr='...')
448
+ return
449
+
450
+ # repr object
451
+ yield from repr_object(node)
452
+
453
+ # Add to seen modules
454
+ # CONTEXT.seen_modules_repr.add(id(node))
455
+ CONTEXT.seen_modules_repr[id(node)] = node
456
+
457
+ try:
458
+ # repr attributes
459
+ yield from repr_attr(node)
460
+ finally:
461
+ if clear_seen:
462
+ CONTEXT.seen_modules_repr = None