loading a pickle file - python-2.7

Question

I'm trying to get a "getting started with pickles" script working. I managed to save a pickle file from a file, and load it. But when I save a pickle file in one file (the main.py in this case) and load it from another, I get an error. I probably missed something small, but can't figure out what.
main.py
import pickle
class Node:
"""This class represents a node"""
def __init__(self, value = None):
self.val = value
def toString(self):
return self.val
class Link:
"""This class represents a link between 2 nodes"""
def __init__(self, sourceNode, targetNode, LinkWigth):
self.source = sourceNode
self.target = targetNode
self.wight = LinkWigth
def setWeight(self, newWeight):
self.wight = newWeight
def toString(self):
return self.wight
class Graph:
"""This class represents a graph"""
def __init__(self):
self.nodes = []
self.links = []
def addNode(self, node):
self.nodes.append(node)
def addLink(self, link):
self.links.append(link)
def getInDegree(self, node):
counter = 0
for link in self.links:
if link.target == node:
counter +=1
else:
print "target is: %s" % link.target.toString()
print "source is: %s" % link.source.toString()
return counter
def toString(self):
for link in self.links:
print link.toString()
for node in self.nodes:
print node.toString()
if __name__ == "__main__":
n1 = Node(4)
l1 = Link(n1, n1, 1)
g = Graph()
g.addNode(n1)
g.addLink(l1)
pickle.dump(g, open('haha', 'wb') )
pickleLoader.py
import pickle
import main
n = main.Node(44)
print n.toString()
g = pickle.load( open('haha', 'rb') )
print "ha"
The error
C:\Users\R\Desktop\pickle test>main.py
C:\Users\R\Desktop\pickle test>pickleLoader.py
44
Traceback (most recent call last):
File "C:\Users\R\Desktop\pickle test\pickleLoader.py", line 7, in <module>
g = pickle.load( open('haha', 'rb') )
File "C:\Program Files\Python27\lib\pickle.py", line 1378, in load
return Unpickler(file).load()
File "C:\Program Files\Python27\lib\pickle.py", line 858, in load
dispatch[key](self)
File "C:\Program Files\Python27\lib\pickle.py", line 1069, in load_inst
klass = self.find_class(module, name)
File "C:\Program Files\Python27\lib\pickle.py", line 1126, in find_class
klass = getattr(mod, name)
AttributeError: 'module' object has no attribute 'Graph'
C:\Users\R\Desktop\pickle test>
I guess that the problem is something with the namespace because main.py has been imported, but I have no idea how to get it working.

Answer 1

This does appear to be related to how the classes are defined in relation to the module. A quick way to allow this to work is to import the components of the main module directly into pickleLoader:
from main import Graph, Node, Link
A better solution might be to move the common components (Graph, Node, Link) into their own module, and then import that module into both main and pickleLoader.