Jason C
Jason C

Reputation: 40396

Problem using `shelve` in "__main__" vs imported module

Consider the following program (Python 3.10), which consists of two files:

mymodule.py

import dataclasses
import shelve
from typing import Optional

# public api
# ----------

@dataclasses.dataclass
class SomeObject:
    some_property: int = 0

def save_object (obj: SomeObject) -> None:
    with shelve.open("theshelf", flag="c") as db:
        db["object"] = obj

def load_object () -> Optional[SomeObject]:
    with shelve.open("theshelf", flag="c") as db:
        return db.get("object")

# tester
# ------

if __name__ == "__main__":
    save_object(SomeObject())
    obj = load_object()

myprogram.py

import mymodule

if __name__ == "__main__":
    # this fails (see description below)
    obj = mymodule.load_object()

There is a module, mymodule.py, which contains an embedded tester program that is executed if you run the module directly. This is a pattern that I commonly use. There is also the "real" program, myprogram.py. This is a silly program but it's representative of the problem I'm having.

The problem is this:

If I run the test code, e.g. python mymodule.py, everything works fine and the shelf file is written.

However, if I first run the test code to serialize the SomeObject and then run the "actual" program (python myprogram.py), I get the following error:

% python mymodule.py 
% python myprogram.py 

Traceback (most recent call last):
  File "/opt/homebrew/Cellar/python@3.12/3.12.2_1/Frameworks/Python.framework/Versions/3.12/lib/python3.12/shelve.py", line 111, in __getitem__
    value = self.cache[key]
            ~~~~~~~~~~^^^^^
KeyError: 'object'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/jason_pro2023/personal/shelf/myprogram.py", line 4, in <module>
    obj = mymodule.load_object()
          ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jason_pro2023/personal/shelf/mymodule.py", line 18, in load_object
    return db.get("object")
           ^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Cellar/python@3.12/3.12.2_1/Frameworks/Python.framework/Versions/3.12/lib/python3.12/shelve.py", line 106, in get
    return self[key]
           ~~~~^^^^^
  File "/opt/homebrew/Cellar/python@3.12/3.12.2_1/Frameworks/Python.framework/Versions/3.12/lib/python3.12/shelve.py", line 114, in __getitem__
    value = Unpickler(f).load()
            ^^^^^^^^^^^^^^^^^^^
AttributeError: Can't get attribute 'SomeObject' on <module '__main__' from '/Users/jason_pro2023/personal/shelf/myprogram.py'>

What I believe is happening is that when I run the tester, the module name is "__main__" and that gets saved in the shelf when the SomeObject is serialized. Then if I run the real program, the module name is "mymodule" (not "__main__"), and so the object can't be deserialized.

My question is: Is there some way I can avoid this but also continue my pattern of embedding test code directly in my modules?

I could write the tester as a separate application, and that's probably what I'll do if I don't find a solution, but if there's a way around this, I'd rather keep the code in the module.

Upvotes: 0

Views: 48

Answers (0)

Related Questions