0
# Type Registry System
1

2
dill provides a flexible type registry system that allows registration of custom types and extension of serialization capabilities to handle new object types that are not supported by default.
3

4
## Core Registry Functions
5

6
### Type Registration
7

8
```python { .api }
9
def register(t):
10
    """
11
    Register a type with the pickler.
12
    
13
    Decorator function that registers a custom pickling function for a specific
14
    type, allowing dill to handle objects of that type during serialization.
15
    
16
    Parameters:
17
    - t: type, the type to register for custom pickling
18
    
19
    Returns:
20
    function: decorator function that takes the pickling function
21
    
22
    Usage:
23
    @dill.register(MyCustomType)
24
    def save_custom_type(pickler, obj):
25
        # Custom pickling logic
26
        pass
27
    """
28

29
def pickle(t, func):
30
    """
31
    Add a type to the pickle dispatch table.
32
    
33
    Directly associates a type with a pickling function in the dispatch table,
34
    enabling automatic handling of objects of that type.
35
    
36
    Parameters:
37
    - t: type, the type to add to dispatch table
38
    - func: function, pickling function that handles objects of type t
39
    
40
    Returns:
41
    None
42
    
43
    Raises:
44
    - TypeError: when type or function parameters are invalid
45
    """
46

47
def extend(use_dill=True):
48
    """
49
    Add or remove dill types to/from the pickle registry.
50
    
51
    Controls whether dill's extended types are available in the standard
52
    pickle module's dispatch table, allowing integration with existing code.
53
    
54
    Parameters:
55
    - use_dill: bool, if True extend dispatch table with dill types,
56
                     if False revert to standard pickle types only
57
    
58
    Returns:
59
    None
60
    """
61
```
62

63
### Registry Inspection
64

65
```python { .api }  
66
def dispatch_table():
67
    """
68
    Get the current dispatch table.
69
    
70
    Returns the current pickle dispatch table showing all registered
71
    types and their associated pickling functions.
72
    
73
    Returns:
74
    dict: mapping of types to pickling functions
75
    """
76
```
77

78
## Usage Examples
79

80
### Basic Type Registration
81

82
```python
83
import dill
84

85
# Define a custom class
86
class CustomData:
87
    def __init__(self, data, metadata=None):
88
        self.data = data
89
        self.metadata = metadata or {}
90
    
91
    def __repr__(self):
92
        return f"CustomData({self.data}, {self.metadata})"
93

94
# Register custom pickling function using decorator
95
@dill.register(CustomData)
96
def save_custom_data(pickler, obj):
97
    """Custom pickling function for CustomData objects."""
98
    # Save the class reference
99
    pickler.save_reduce(_restore_custom_data, 
100
                       (obj.data, obj.metadata), obj=obj)
101

102
def _restore_custom_data(data, metadata):
103
    """Helper function to restore CustomData objects."""
104
    return CustomData(data, metadata)
105

106
# Test the registration
107
custom_obj = CustomData([1, 2, 3], {"version": "1.0"})
108

109
# Now it can be pickled
110
serialized = dill.dumps(custom_obj)
111
restored = dill.loads(serialized)
112

113
print(f"Original: {custom_obj}")
114
print(f"Restored: {restored}")
115
```
116

117
### Direct Dispatch Table Modification
118

119
```python
120
import dill
121

122
class SpecialContainer:
123
    def __init__(self, items):
124
        self.items = list(items)
125
    
126
    def __eq__(self, other):
127
        return isinstance(other, SpecialContainer) and self.items == other.items
128

129
def pickle_special_container(pickler, obj):
130
    """Pickling function for SpecialContainer."""
131
    # Use pickle protocol to save class and items
132
    pickler.save_reduce(SpecialContainer, (obj.items,), obj=obj)
133

134
# Add to dispatch table directly
135
dill.pickle(SpecialContainer, pickle_special_container)
136

137
# Test the registration
138
container = SpecialContainer([1, 2, 3, 4])
139
pickled_data = dill.dumps(container)
140
unpickled_container = dill.loads(pickled_data)
141

142
print(f"Original == Restored: {container == unpickled_container}")
143
```
144

145
### Complex Type Registration
146

147
```python
148
import dill
149
import weakref
150

151
class Node:
152
    """Tree node with parent/child relationships."""
153
    def __init__(self, value, parent=None):
154
        self.value = value
155
        self.children = []
156
        self._parent_ref = None
157
        if parent:
158
            self.set_parent(parent)
159
    
160
    def set_parent(self, parent):
161
        if self._parent_ref:
162
            old_parent = self._parent_ref()
163
            if old_parent:
164
                old_parent.children.remove(self)
165
        
166
        self._parent_ref = weakref.ref(parent) if parent else None
167
        if parent:
168
            parent.children.append(self)
169
    
170
    @property
171
    def parent(self):
172
        return self._parent_ref() if self._parent_ref else None
173

174
# Register complex type with custom handling
175
@dill.register(Node)
176
def save_node(pickler, obj):
177
    """Custom pickling for Node objects."""
178
    # Save value and children (parent will be reconstructed)
179
    pickler.save_reduce(_restore_node, 
180
                       (obj.value, obj.children), obj=obj)
181

182
def _restore_node(value, children):
183
    """Restore Node with proper parent/child relationships."""
184
    node = Node(value)
185
    
186
    # Restore children and set parent relationships
187
    for child in children:
188
        if isinstance(child, Node):
189
            child.set_parent(node)
190
        else:
191
            # Handle case where child was also pickled/unpickled
192
            node.children.append(child)
193
    
194
    return node
195

196
# Test complex type registration
197
root = Node("root")
198
child1 = Node("child1", root)
199
child2 = Node("child2", root)
200
grandchild = Node("grandchild", child1)
201

202
# Pickle the entire tree
203
tree_data = dill.dumps(root)
204
restored_root = dill.loads(tree_data)
205

206
print(f"Root children: {[child.value for child in restored_root.children]}")
207
print(f"Child1 parent: {restored_root.children[0].parent.value}")
208
```
209

210
### Registry Management
211

212
```python
213
import dill
214

215
class RegistryManager:
216
    """Manage custom type registrations."""
217
    
218
    def __init__(self):
219
        self.registered_types = {}
220
        self.original_dispatch = None
221
    
222
    def register_type(self, type_class, save_func, restore_func=None):
223
        """Register a type with save and optional restore functions."""
224
        if restore_func:
225
            # Store restore function for later use
226
            self.registered_types[type_class] = {
227
                'save': save_func,
228
                'restore': restore_func
229
            }
230
            
231
            # Create wrapper that includes restore function
232
            def save_wrapper(pickler, obj):
233
                pickler.save_reduce(restore_func, save_func(obj), obj=obj)
234
            
235
            dill.register(type_class)(save_wrapper)
236
        else:
237
            # Direct registration
238
            dill.register(type_class)(save_func)
239
            self.registered_types[type_class] = {'save': save_func}
240
    
241
    def unregister_type(self, type_class):
242
        """Remove type from registry."""
243
        if type_class in self.registered_types:
244
            # Remove from dill's dispatch table
245
            dispatch = dill.dispatch_table()
246
            if type_class in dispatch:
247
                del dispatch[type_class]
248
            
249
            del self.registered_types[type_class]
250
    
251
    def list_registered_types(self):
252
        """List all registered custom types."""
253
        return list(self.registered_types.keys())
254
    
255
    def backup_dispatch_table(self):
256
        """Backup current dispatch table."""
257
        self.original_dispatch = dill.dispatch_table().copy()
258
    
259
    def restore_dispatch_table(self):
260
        \"\"\"Restore original dispatch table.\"\"\"
261
        if self.original_dispatch:
262
            current_dispatch = dill.dispatch_table()
263
            current_dispatch.clear()
264
            current_dispatch.update(self.original_dispatch)
265

266
# Usage example
267
registry = RegistryManager()
268

269
# Backup original state
270
registry.backup_dispatch_table()
271

272
# Register multiple types
273
registry.register_type(
274
    CustomData,
275
    lambda obj: (obj.data, obj.metadata),
276
    lambda data, meta: CustomData(data, meta)
277
)
278

279
registry.register_type(
280
    SpecialContainer,
281
    lambda obj: (obj.items,),
282
    lambda items: SpecialContainer(items)
283
)
284

285
print(f"Registered types: {[t.__name__ for t in registry.list_registered_types()]}")
286

287
# Test registrations
288
test_objects = [
289
    CustomData([1, 2, 3], {"test": True}),
290
    SpecialContainer(["a", "b", "c"])
291
]
292

293
for obj in test_objects:
294
    try:
295
        data = dill.dumps(obj)
296
        restored = dill.loads(data)
297
        print(f"✓ {type(obj).__name__}: Successfully pickled and restored")
298
    except Exception as e:
299
        print(f"✗ {type(obj).__name__}: Failed - {e}")
300

301
# Clean up
302
registry.restore_dispatch_table()
303
```
304

305
## Advanced Registry Techniques
306

307
### Conditional Type Registration
308

309
```python
310
import dill
311
import sys
312

313
class ConditionalRegistry:
314
    """Register types based on conditions."""
315
    
316
    @staticmethod
317
    def register_if_available(type_name, module_name, save_func, restore_func=None):
318
        """Register type only if module is available."""
319
        try:
320
            module = __import__(module_name)
321
            type_class = getattr(module, type_name)
322
            
323
            if restore_func:
324
                def save_wrapper(pickler, obj):
325
                    pickler.save_reduce(restore_func, save_func(obj), obj=obj)
326
                dill.register(type_class)(save_wrapper)
327
            else:
328
                dill.register(type_class)(save_func)
329
            
330
            print(f"✓ Registered {type_name} from {module_name}")
331
            return True
332
            
333
        except (ImportError, AttributeError) as e:
334
            print(f"✗ Could not register {type_name}: {e}")
335
            return False
336
    
337
    @staticmethod
338
    def register_for_python_version(min_version, type_class, save_func):
339
        """Register type only for specific Python versions."""
340
        if sys.version_info >= min_version:
341
            dill.register(type_class)(save_func)
342
            print(f"✓ Registered {type_class.__name__} for Python {sys.version_info[:2]}")
343
            return True
344
        else:
345
            print(f"✗ Skipped {type_class.__name__} - requires Python {min_version}")
346
            return False
347

348
# Example usage
349
registry = ConditionalRegistry()
350

351
# Register numpy array if numpy is available
352
registry.register_if_available(
353
    'ndarray', 'numpy',
354
    lambda obj: (obj.tobytes(), obj.dtype, obj.shape),
355
    lambda data, dtype, shape: __import__('numpy').frombuffer(data, dtype).reshape(shape)
356
)
357

358
# Register pathlib.Path for Python 3.4+
359
import pathlib
360
registry.register_for_python_version(
361
    (3, 4),
362
    pathlib.Path,
363
    lambda obj: str(obj)
364
)
365
```
366

367
### Dynamic Type Discovery
368

369
```python
370
import dill
371
import inspect
372

373
class AutoRegistry:
374
    """Automatically discover and register types."""
375
    
376
    def __init__(self):
377
        self.auto_registered = set()
378
    
379
    def auto_register_module(self, module, save_pattern=None, exclude=None):
380
        """Automatically register types from a module."""
381
        exclude = exclude or []
382
        registered_count = 0
383
        
384
        for name, obj in inspect.getmembers(module):
385
            # Skip private and excluded items
386
            if name.startswith('_') or name in exclude:
387
                continue
388
            
389
            # Only register classes
390
            if inspect.isclass(obj):
391
                # Check if we can create a default save function
392
                if self._can_auto_register(obj):
393
                    self._auto_register_type(obj)
394
                    registered_count += 1
395
        
396
        print(f\"Auto-registered {registered_count} types from {module.__name__}\")\n        return registered_count\n    \n    def _can_auto_register(self, type_class):\n        \"\"\"Check if type can be auto-registered.\"\"\"\n        try:\n            # Check if type has __dict__ or __slots__\n            if hasattr(type_class, '__dict__') or hasattr(type_class, '__slots__'):\n                return True\n            \n            # Check if it's a simple dataclass-like structure\n            if hasattr(type_class, '__init__'):\n                sig = inspect.signature(type_class.__init__)\n                # Simple heuristic: if all parameters have defaults or annotations\n                return len(sig.parameters) <= 5  # Reasonable limit\n        \n        except Exception:\n            pass\n        \n        return False\n    \n    def _auto_register_type(self, type_class):\n        \"\"\"Automatically register a type with default behavior.\"\"\"\n        if type_class in self.auto_registered:\n            return\n        \n        def auto_save_func(pickler, obj):\n            # Generic save function using __dict__ or __getstate__\n            if hasattr(obj, '__getstate__'):\n                state = obj.__getstate__()\n                pickler.save_reduce(_auto_restore_with_state, \n                                  (type_class, state), obj=obj)\n            elif hasattr(obj, '__dict__'):\n                pickler.save_reduce(_auto_restore_with_dict,\n                                  (type_class, obj.__dict__), obj=obj)\n            else:\n                # Fallback to trying __reduce__\n                pickler.save_reduce(type_class, (), obj=obj)\n        \n        dill.register(type_class)(auto_save_func)\n        self.auto_registered.add(type_class)\n        print(f\"Auto-registered {type_class.__name__}\")\n\ndef _auto_restore_with_state(type_class, state):\n    \"\"\"Restore object using __setstate__.\"\"\"\n    obj = type_class.__new__(type_class)\n    if hasattr(obj, '__setstate__'):\n        obj.__setstate__(state)\n    else:\n        obj.__dict__.update(state)\n    return obj\n\ndef _auto_restore_with_dict(type_class, obj_dict):\n    \"\"\"Restore object using __dict__.\"\"\"\n    obj = type_class.__new__(type_class)\n    obj.__dict__.update(obj_dict)\n    return obj\n\n# Example usage\nauto_registry = AutoRegistry()\n\n# Auto-register types from a custom module\n# auto_registry.auto_register_module(my_custom_module, exclude=['BaseClass'])\n```\n\n## Integration with Existing Code\n\n### Backward Compatibility\n\n```python\nimport dill\nimport pickle\n\nclass CompatibilityManager:\n    \"\"\"Manage compatibility between dill and standard pickle.\"\"\"\n    \n    def __init__(self):\n        self.dill_enabled = True\n    \n    def enable_dill_types(self):\n        \"\"\"Enable dill types in pickle dispatch table.\"\"\"\n        dill.extend(use_dill=True)\n        self.dill_enabled = True\n        print(\"Dill types enabled in pickle\")\n    \n    def disable_dill_types(self):\n        \"\"\"Disable dill types, use only standard pickle.\"\"\"\n        dill.extend(use_dill=False)\n        self.dill_enabled = False\n        print(\"Dill types disabled, using standard pickle only\")\n    \n    def test_compatibility(self, test_objects):\n        \"\"\"Test objects with both pickle and dill.\"\"\"\n        results = {}\n        \n        for name, obj in test_objects.items():\n            results[name] = {\n                'dill': self._test_with_dill(obj),\n                'pickle': self._test_with_pickle(obj)\n            }\n        \n        return results\n    \n    def _test_with_dill(self, obj):\n        try:\n            data = dill.dumps(obj)\n            dill.loads(data)\n            return \"✓ Success\"\n        except Exception as e:\n            return f\"✗ Failed: {e}\"\n    \n    def _test_with_pickle(self, obj):\n        try:\n            data = pickle.dumps(obj)\n            pickle.loads(data)\n            return \"✓ Success\"\n        except Exception as e:\n            return f\"✗ Failed: {e}\"\n\n# Usage\ncompat = CompatibilityManager()\n\n# Test various objects\ntest_objects = {\n    'function': lambda x: x + 1,\n    'class': CustomData,\n    'instance': CustomData([1, 2, 3]),\n    'list': [1, 2, 3, 4, 5]\n}\n\nprint(\"Testing with dill types enabled:\")\ncompat.enable_dill_types()\nresults_enabled = compat.test_compatibility(test_objects)\n\nprint(\"\\nTesting with dill types disabled:\")\ncompat.disable_dill_types()\nresults_disabled = compat.test_compatibility(test_objects)\n\n# Compare results\nprint(\"\\nCompatibility Results:\")\nprint(\"-\" * 50)\nfor name in test_objects:\n    print(f\"{name}:\")\n    print(f\"  Dill enabled  - pickle: {results_enabled[name]['pickle']}\")\n    print(f\"  Dill disabled - pickle: {results_disabled[name]['pickle']}\")\n    print(f\"  Dill enabled  - dill:   {results_enabled[name]['dill']}\")\n\n# Re-enable for normal operation\ncompat.enable_dill_types()\n```\n\n## Best Practices\n\n### Registry Guidelines\n\n1. **Performance**: Keep pickling functions simple and fast\n2. **Robustness**: Handle edge cases and provide error recovery\n3. **Compatibility**: Test with different Python versions and environments\n4. **Documentation**: Document custom types and their pickling behavior\n5. **Testing**: Verify that pickled objects restore correctly with full functionality\n\n### Error Handling\n\n```python\nimport dill\nimport logging\n\nclass RobustRegistry:\n    \"\"\"Registry with comprehensive error handling.\"\"\"\n    \n    def __init__(self):\n        self.logger = logging.getLogger(__name__)\n    \n    def safe_register(self, type_class, save_func, restore_func=None):\n        \"\"\"Register type with error handling.\"\"\"\n        try:\n            # Validate inputs\n            if not isinstance(type_class, type):\n                raise TypeError(\"type_class must be a type\")\n            \n            if not callable(save_func):\n                raise TypeError(\"save_func must be callable\")\n            \n            if restore_func and not callable(restore_func):\n                raise TypeError(\"restore_func must be callable\")\n            \n            # Test the registration with a dummy object\n            self._test_registration(type_class, save_func, restore_func)\n            \n            # Perform actual registration\n            if restore_func:\n                def wrapper(pickler, obj):\n                    try:\n                        args = save_func(obj)\n                        pickler.save_reduce(restore_func, args, obj=obj)\n                    except Exception as e:\n                        self.logger.error(f\"Error saving {type_class.__name__}: {e}\")\n                        raise\n                \n                dill.register(type_class)(wrapper)\n            else:\n                dill.register(type_class)(save_func)\n            \n            self.logger.info(f\"Successfully registered {type_class.__name__}\")\n            return True\n            \n        except Exception as e:\n            self.logger.error(f\"Failed to register {type_class.__name__}: {e}\")\n            return False\n    \n    def _test_registration(self, type_class, save_func, restore_func):\n        \"\"\"Test registration with dummy data.\"\"\"\n        # This would need to be implemented based on specific type requirements\n        pass\n\n# Usage\nrobust_registry = RobustRegistry()\nlogging.basicConfig(level=logging.INFO)\n\nsuccess = robust_registry.safe_register(\n    CustomData,\n    lambda obj: (obj.data, obj.metadata),\n    lambda data, meta: CustomData(data, meta)\n)\n\nif success:\n    print(\"Registration successful\")\nelse:\n    print(\"Registration failed - check logs\")\n```