Can't pickle then unpickle a tree sequence: "no attribute '_table_metadata_schemas'"

[hyanwong]

ts = msprime.simulate(10)
with open('ts.pickle', 'wb') as handle:
    pickle.dump(ts, handle, protocol=pickle.HIGHEST_PROTOCOL)
print(ts.draw_text())
with open('ts.pickle', 'rb') as handle:
    ts = pickle.load(handle)
print(ts.draw_text())

Gives

AttributeError                            Traceback (most recent call last)
<ipython-input-52-2b8d7332ec21> in <module>
      7 with open('ts.pickle', 'rb') as handle:
      8     ts = pickle.load(handle)
----> 9 print(ts.draw_text())
     10 
     11 

~/anaconda3/lib/python3.9/site-packages/tskit/trees.py in draw_text(self, **kwargs)
   5427     def draw_text(self, **kwargs):
   5428         # TODO document this method.
-> 5429         return str(drawing.TextTreeSequence(self, **kwargs))
   5430 
   5431     ############################################

~/anaconda3/lib/python3.9/site-packages/tskit/drawing.py in __init__(self, ts, node_labels, use_ascii, time_label_format, position_label_format, order)
   1250             position_label_format.format(x) for x in ts.breakpoints()
   1251         ]
-> 1252         trees = [
   1253             VerticalTextTree(
   1254                 tree,

~/anaconda3/lib/python3.9/site-packages/tskit/drawing.py in <listcomp>(.0)
   1251         ]
   1252         trees = [
-> 1253             VerticalTextTree(
   1254                 tree,
   1255                 max_tree_height="ts",

~/anaconda3/lib/python3.9/site-packages/tskit/drawing.py in __init__(self, tree, node_labels, max_tree_height, use_ascii, orientation, order)
   1455                 self.node_labels[node] = label
   1456 
-> 1457         self._assign_time_positions()
   1458         self._assign_traversal_positions()
   1459         self.canvas = np.zeros((self.height, self.width), dtype=str)

~/anaconda3/lib/python3.9/site-packages/tskit/drawing.py in _assign_time_positions(self)
   1481         # account here. Presumably we need to get the maximum number of mutations
   1482         # per branch.
-> 1483         self.time_position, total_depth = node_time_depth(
   1484             tree, max_tree_height=self.max_tree_height
   1485         )

~/anaconda3/lib/python3.9/site-packages/tskit/drawing.py in node_time_depth(tree, min_branch_length, max_tree_height)
   1385         assert max_tree_height == "ts"
   1386         ts = tree.tree_sequence
-> 1387         for node in ts.nodes():
   1388             time_node_map[node.time].append(node.id)
   1389         node_edges = collections.defaultdict(list)

~/anaconda3/lib/python3.9/site-packages/tskit/trees.py in __getitem__(self, index)
   3234         if index < 0 or index >= len(self):
   3235             raise IndexError("Index out of bounds")
-> 3236         return self.getter(index)
   3237 
   3238 

~/anaconda3/lib/python3.9/site-packages/tskit/trees.py in node(self, id_)
   4516             individual=individual,
   4517             encoded_metadata=metadata,
-> 4518             metadata_decoder=self.table_metadata_schemas.node.decode_row,
   4519         )
   4520 

~/anaconda3/lib/python3.9/site-packages/tskit/trees.py in table_metadata_schemas(self)
   3654         The set of metadata schemas for the tables in this tree sequence.
   3655         """
-> 3656         return self._table_metadata_schemas
   3657 
   3658     @property

AttributeError: 'TreeSequence' object has no attribute '_table_metadata_schemas'

[benjeffery]

Weird, we're clearly missing a test! Thanks for the report.

[jeromekelleher]

Ah - we tested the tables for pickleability, but apparently not the tree sequence. Good catch @hyanwong.

We should try the Tree class as well, and see if we get a comprehensible error message when trying to pickle it. We probably don't want to actually support pickling I think, at least not until we have random access for trees.

[benjeffery]

We had tests for tree sequence pickling, but they only checked the table content of the tree sequence, not that the class was working properly.