Better dict for configuration

When I use dict as a config of deep learning algorithms, I often run multiple training runs, varying the config value.

As a trivial example:

config = {
  'lr': 1e-3
}

And then I run multiple iterations varying the ‘lr’.

When doing that, I feel worried about making mistake like updating the dict with wrong key, e.g.,

config['Lr'] = 1e-4  # This is not the key I wanted to update.

To address such concerns, I ended up writing a dict that doesn’t allow new keys. I call it as FrozenKeyDict. See the test first. Note that it raises an exception when ‘c’ is being set.

import unittest

from util.frozen_key_dict import FrozenKeyDict


class FrozenKeyDicTestCase(unittest.TestCase):

    def test_missing_key(self):
        d = FrozenKeyDict({"a": 1, "b": 2})
        self.assertEqual(d["a"], 1)
        d["a"] = 2
        self.assertEqual(d["a"], 2)
        with self.assertRaises(AssertionError):
            d["c"] = 3


if __name__ == "__main__":
    unittest.main()

Here’s the implementation.

from collections import UserDict
from typing import Any


class FrozenKeyDict(UserDict):

    def __init__(self, *args, **kwargs):
        self._frozen = False
        super().__init__(*args, **kwargs)
        self._frozen = True

    def __setitem__(self, key: Any, item: Any) -> None:
        if self._frozen:
            assert key in self, f"Key {key=} not in {self.keys()}"
        return super().__setitem__(key, item)

Copy paste the code and use it freely.

Passion is like genius; a miracle.

Comments

Leave a Reply