Why am I getting a KeyError when trying to use an Enum as a dictionary key in another file?

Question

I have two files, enum_test.py and print_test_file.py. I defined an Enum class in enum_test.py and created a dictionary using members of that Enum class as keys. When I tried to use that dictionary in print_test_file.py I got a key error.

Here is the code for enum_test.py:

from enum import Enum
import print_test_file

class MyEnum(Enum):
    A = 1
    B = 2

def main():
    enumDict = dict()
    enumDict[MyEnum.A] = 'abcd'
    enumDict[MyEnum.B] = 'efgh'

    print(enumDict[MyEnum.A])
    print_test_file.print_test(enumDict)

if __name__ == "__main__":
    main()

and here is the code from print_test_file.py

import enum_test

def print_test(enumDict):
    print(enumDict[enum_test.MyEnum.A])

I was expecting that running enum_test would produce the following output:

abcd
abcd

but instead it only prints the first abcd and then throws KeyError: <MyEnum.A: 1>

How can I use the dictionary I created in enum_test.py from print_test_file.py?

Answer 1

This is exactly why you should never import the same module that you're running as a script. You end up with two completely independent module objects, enum_test and __main__, with their own completely independent global namespaces, containing separate objects that were constructed from the same code.

__main__.main() builds a dict, and fills it with __main__.MyEnum keys.

print_test_file.print_test takes that dict as a parameter. But then it tries to search it using enum_test.MyEnum.A. And that key does not exist in that dict. So you get a KeyError.

---

Your modules also have circular dependencies, which is obscured by this problem. I'm pretty sure you'd actually get away with the circular dependencies in this case if nothing else was wrong, but it's still confusing to have to think through the order of how top-level module code gets executed, so it's better to avoid them.

---

The easy way to fix both problems at once is to move the shared code into a separate shared module, which both your script and your test module can import.

# enum_test.py
from my_enum import MyEnum
import print_test_file

def main():
    enumDict = dict()
    enumDict[MyEnum.A] = 'abcd'
    enumDict[MyEnum.B] = 'efgh'

    print(enumDict[MyEnum.A])
    print_test_file.print_test(enumDict)

if __name__ == "__main__":
    main()

---

# print_test_file.py
import my_enum

def print_test(enumDict):
    print(enumDict[my_enum.MyEnum.A])

---

# my_enum.py
from enum import Enum

class MyEnum(Enum):
    A = 1
    B = 2

Now, there's no module that imports the script, and no module that imports the module that imports it. (In technical terms, instead of a graph with cycles in it, you just have a tree, with the script at the root.)

---

Every few years, someone suggests changing Python to eliminate this problem, but there really is no good answer. They could make it an error to import the same module being run as a script, but that could break some uncommon but important use cases (like multiprocessing). Or they could make it “just work” by doing an implicit sys.modules['enum_test'] = sys.modules['__main__'] somewhere—but there really is no “somewhere” that wouldn’t turn the circular dependency in your code (and almost every other example of this problem) from mostly harmless to a serious bug.

(By the way, I vaguely remember that in one of the discussions, one of the core devs who’s also a teacher mentioned that very few of his students ever run into this problem, but the ones who do tend to end up near the top of the class, so at least you have that to feel good about.)