2
2
3
3
import logging
4
4
import types
5
+ from collections import UserDict
5
6
from typing import TYPE_CHECKING , TypeVar
6
7
7
8
from .._abc import Instrument
@@ -21,7 +22,7 @@ def _public(fn: T) -> T:
21
22
return fn
22
23
23
24
24
- class Instruments (dict [str , dict [Instrument , None ]]):
25
+ class Instruments (UserDict [str , dict [Instrument , None ]]):
25
26
"""A collection of `trio.abc.Instrument` organized by hook.
26
27
27
28
Instrumentation calls are rather expensive, and we don't want a
@@ -34,7 +35,7 @@ class Instruments(dict[str, dict[Instrument, None]]):
34
35
__slots__ = ()
35
36
36
37
def __init__ (self , incoming : Sequence [Instrument ]) -> None :
37
- self [ "_all" ] = {}
38
+ super (). __init__ ({ "_all" : {}})
38
39
for instrument in incoming :
39
40
self .add_instrument (instrument )
40
41
@@ -48,9 +49,9 @@ def add_instrument(self, instrument: Instrument) -> None:
48
49
If ``instrument`` is already active, does nothing.
49
50
50
51
"""
51
- if instrument in self ["_all" ]:
52
+ if instrument in self . data ["_all" ]:
52
53
return
53
- self ["_all" ][instrument ] = None
54
+ self . data ["_all" ][instrument ] = None
54
55
try :
55
56
for name in dir (instrument ):
56
57
if name .startswith ("_" ):
@@ -63,7 +64,7 @@ def add_instrument(self, instrument: Instrument) -> None:
63
64
if isinstance (impl , types .MethodType ) and impl .__func__ is prototype :
64
65
# Inherited unchanged from _abc.Instrument
65
66
continue
66
- self .setdefault (name , {})[instrument ] = None
67
+ self .data . setdefault (name , {})[instrument ] = None
67
68
except :
68
69
self .remove_instrument (instrument )
69
70
raise
@@ -83,12 +84,12 @@ def remove_instrument(self, instrument: Instrument) -> None:
83
84
84
85
"""
85
86
# If instrument isn't present, the KeyError propagates out
86
- self ["_all" ].pop (instrument )
87
- for hookname , instruments in list (self .items ()):
87
+ self . data ["_all" ].pop (instrument )
88
+ for hookname , instruments in list (self .data . items ()):
88
89
if instrument in instruments :
89
90
del instruments [instrument ]
90
91
if not instruments :
91
- del self [hookname ]
92
+ del self . data [hookname ]
92
93
93
94
def call (
94
95
self ,
@@ -103,7 +104,7 @@ def call(
103
104
if "before_task_step" in instruments:
104
105
instruments.call("before_task_step", task)
105
106
"""
106
- for instrument in list (self [hookname ]):
107
+ for instrument in list (self . data [hookname ]):
107
108
try :
108
109
getattr (instrument , hookname )(* args )
109
110
except BaseException :
0 commit comments