Skip to content

Instantly share code, notes, and snippets.

@attentive
Last active October 15, 2022 05:18
Show Gist options
  • Save attentive/f56f2141104fe0e13fbc4b9ae5d0504f to your computer and use it in GitHub Desktop.
Save attentive/f56f2141104fe0e13fbc4b9ae5d0504f to your computer and use it in GitHub Desktop.
Versions of Python dict and collections.defaultdict that allow one to use tuples as keys in a natural way
# -*- coding: utf-8 -*-
from collections import defaultdict
# A defaultdict variant of TupleDict with some improvements (eg better exceptions)
class DefaultTupleDict(defaultdict):
__slots__ = ()
@classmethod
def fromkeys(cls, keys, v=None):
return defaultdict.fromkeys(keys, v)
def __repr__(self):
return defaultdict.__repr__(self)
def __str__(self):
return repr(self)
def __init__(self, mapping=(), **kwargs):
defaultdict.__init__(self, DefaultTupleDict, mapping, **kwargs)
def _visit_tuple_key(self, key: tuple, cb):
node = self
(k, *ks) = key
while ks:
if node is not None:
if hasattr(node, '__getitem__'):
node = defaultdict.__getitem__(node, k)
else:
raise TypeError(
f"'{type(node)}' object is not subscriptable")
(k, *ks) = ks
else:
raise TypeError(
f"'{type(node)}' object is not subscriptable")
return cb(node, k)
@staticmethod
def _check_caller(obj, method):
try:
if hasattr(obj, method):
caller = getattr(obj, method)
if callable(caller):
return caller
raise TypeError(
f"'{type(obj)}' object has no method '{method}'")
except:
raise TypeError(f"'{type(obj)}' object has no method '{method}'")
def _visit_key(self, method, key, *args, **kwargs):
if isinstance(key, tuple):
return self._visit_tuple_key(key, lambda node, k: DefaultTupleDict._check_caller(node, method)(k, *args, **kwargs))
else:
return DefaultTupleDict._check_caller(defaultdict, method)(self, key, *args, **kwargs)
def __contains__(self, key):
return self._visit_key('__contains__', key)
def __getitem__(self, key):
return self._visit_key('__getitem__', key)
def __setitem__(self, key, value):
self._visit_key('__setitem__', key, value)
def __delitem__(self, key):
self._visit_key('__delitem__', key)
def get(self, *keys, default=None):
try:
return self._visit_key('__getitem__', tuple(keys))
except KeyError:
return default
def setdefault(self, *keys, default=None):
return self._visit_key('setdefault', tuple(keys), default)
def pop(self, *keys, default=None):
return self._visit_key('pop', tuple(keys), default)
def update(self, mapping=(), **kwargs):
defaultdict.update(self, mapping, **kwargs)
def copy(self):
return type(self)(self)
# foo = DefaultTupleDict()
# print(f"foo[1,2] = {foo[1,2]}")
# print(f"foo[1,2,3] = {foo[1,2,3]}")
# print(f"foo[1,3,3] = {foo[1,3,3]}")
# foo[1,2] = 3
# print(f"foo[1,2] = {foo[1,2]}")
# foo[1,2,3]
# -*- coding: utf-8 -*
# Inspired by https://stackoverflow.com/questions/3387691/how-to-perfectly-override-a-dict
class TupleDict(dict):
__slots__ = ()
@classmethod
def fromkeys(cls, keys, v=None):
return dict.fromkeys(keys, v)
def __repr__(self):
return dict.__repr__(self)
def __str__(self):
return repr(self)
def __init__(self, mapping=(), **kwargs):
dict.__init__(self, mapping, **kwargs)
def _visit_tuple_key(self, key: tuple, cb):
node = self
(k, *ks) = key
while ks:
if node is not None:
if hasattr(node, '__getitem__'):
node = dict.__getitem__(node, k)
else:
raise TypeError(
f"'{type(node)}' object is not subscriptable")
(k, *ks) = ks
else:
raise TypeError(
f"'{type(node)}' object is not subscriptable")
return cb(node, k)
@staticmethod
def _check_caller(obj, method):
try:
if hasattr(obj, method):
caller = getattr(obj, method)
if callable(caller):
return caller
raise TypeError(
f"'{type(obj)}' object has no method '{method}'")
except:
raise TypeError(f"'{type(obj)}' object has no method '{method}'")
def _visit_key(self, method, key, *args, **kwargs):
if isinstance(key, tuple):
return self._visit_tuple_key(key, lambda node, k: TupleDict._check_caller(node, method)(k, *args, **kwargs))
else:
return TupleDict._check_caller(dict, method)(self, key, *args, **kwargs)
def __contains__(self, key):
return self._visit_key('__contains__', key)
def __getitem__(self, key):
return self._visit_key('__getitem__', key)
def __setitem__(self, key, value):
self._visit_key('__setitem__', key, value)
def __delitem__(self, key):
self._visit_key('__delitem__', key)
def get(self, *keys, default=None):
try:
return self._visit_key('__getitem__', tuple(keys))
except KeyError:
return default
def setdefault(self, *keys, default=None):
return self._visit_key('setdefault', tuple(keys), default)
def pop(self, *keys, default=None):
return self._visit_key('pop', tuple(keys), default)
def update(self, mapping=(), **kwargs):
dict.update(self, mapping, **kwargs)
def copy(self):
return type(self)(self)
# # This subclass of dict allows a nice usage when keys are tuples, eg:
# #
# foo = TupleDict()
# foo[1] = TupleDict()
# foo[1][2] = TupleDict()
# foo[1][2][3] = 4
# foo[1, 2, 3]
# # 4
# foo[(1, 2, 3)]
# # 4
# foo[1, 2, 3] = 5
# foo[1, 2, 3]
# # 5
# #
# # To set it up where you didn't need to specify TupleDict() the whole time, you
# # can subclass dict instead with the same approach: an exercise.
@attentive
Copy link
Author

Seem to work fairly well and I might actually try using these in a work setting, but very little testing carried out.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment