Fix pickling SmartLists (fixes #289)

CHANGELOG

@@ -4,6 +4,7 @@ v0.6.5 (unreleased):
 - Added support for Python 3.11.
 - Fixed parsing of leading zeros in named HTML entities. (#288)
 - Fixed memory leak parsing tags. (#303)
+- Fixed pickling SmartList objects. (#289)
 
 v0.6.4 (released February 14, 2022):
 

docs/changelog.rst

@@ -13,6 +13,8 @@ Unreleased
   (`#288 <https://github.com/earwig/mwparserfromhell/issues/288>`_)
 - Fixed memory leak parsing tags.
   (`#303 <https://github.com/earwig/mwparserfromhell/issues/303>`_)
+- Fixed pickling SmartList objects.
+  (`#289 <https://github.com/earwig/mwparserfromhell/issues/289>`_)
 
 v0.6.4
 ------

src/mwparserfromhell/smart_list/list_proxy.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2012-2020 Ben Kurtovic <ben.kurtovic@gmail.com>
+# Copyright (C) 2012-2023 Ben Kurtovic <ben.kurtovic@gmail.com>
 # Copyright (C) 2019-2020 Yuri Astrakhan <YuriAstrakhan@gmail.com>
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -19,6 +19,8 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import weakref
+
 from .utils import _SliceNormalizerMixIn, inheritdoc
 
 
@@ -30,11 +32,21 @@ class ListProxy(_SliceNormalizerMixIn, list):
     it builds it dynamically using the :meth:`_render` method.
     """
 
+    __slots__ = ("__weakref__", "_parent", "_sliceinfo")
+
     def __init__(self, parent, sliceinfo):
         super().__init__()
         self._parent = parent
         self._sliceinfo = sliceinfo
 
+    def __reduce_ex__(self, protocol: int) -> tuple:
+        return (ListProxy, (self._parent, self._sliceinfo), ())
+
+    def __setstate__(self, state: tuple) -> None:
+        # Reregister with the parent
+        child_ref = weakref.ref(self, self._parent._delete_child)
+        self._parent._children[id(child_ref)] = (child_ref, self._sliceinfo)
+
     def __repr__(self):
         return repr(self._render())
 

src/mwparserfromhell/smart_list/smart_list.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2012-2020 Ben Kurtovic <ben.kurtovic@gmail.com>
+# Copyright (C) 2012-2023 Ben Kurtovic <ben.kurtovic@gmail.com>
 # Copyright (C) 2019-2020 Yuri Astrakhan <YuriAstrakhan@gmail.com>
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -49,12 +49,16 @@ class SmartList(_SliceNormalizerMixIn, list):
         [0, 1, 2, 3, 4]
     """
 
-    def __init__(self, iterable=None):
-        if iterable:
-            super().__init__(iterable)
-        else:
-            super().__init__()
-        self._children = {}
+    __slots__ = ("_children",)
+
+    def __new__(cls, *args, **kwargs):
+        obj = super().__new__(cls, *args, **kwargs)
+        obj._children = {}
+        return obj
+
+    def __reduce_ex__(self, protocol: int) -> tuple:
+        # Detach children when pickling
+        return (SmartList, (), None, iter(self))
 
     def __getitem__(self, key):
         if not isinstance(key, slice):

src/mwparserfromhell/smart_list/utils.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2012-2016 Ben Kurtovic <ben.kurtovic@gmail.com>
+# Copyright (C) 2012-2023 Ben Kurtovic <ben.kurtovic@gmail.com>
 # Copyright (C) 2019-2020 Yuri Astrakhan <YuriAstrakhan@gmail.com>
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -37,6 +37,8 @@ def inheritdoc(method):
 class _SliceNormalizerMixIn:
     """MixIn that provides a private method to normalize slices."""
 
+    __slots__ = ()
+
     def _normalize_slice(self, key, clamp=False):
         """Return a slice equivalent to the input *key*, standardized."""
         if key.start is None:

src/mwparserfromhell/utils.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2012-2020 Ben Kurtovic <ben.kurtovic@gmail.com>
+# Copyright (C) 2012-2023 Ben Kurtovic <ben.kurtovic@gmail.com>
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -23,10 +23,20 @@ This module contains accessory functions for other parts of the library. Parser
 users generally won't need stuff from here.
 """
 
+from __future__ import annotations
+
 __all__ = ["parse_anything"]
 
+import typing
+from typing import Any
+
+if typing.TYPE_CHECKING:
+    from .wikicode import Wikicode
+
 
-def parse_anything(value, context=0, skip_style_tags=False):
+def parse_anything(
+    value: Any, context: int = 0, *, skip_style_tags: bool = False
+) -> Wikicode:
     """Return a :class:`.Wikicode` for *value*, allowing multiple types.
 
     This differs from :meth:`.Parser.parse` in that we accept more than just a
@@ -58,11 +68,13 @@ def parse_anything(value, context=0, skip_style_tags=False):
     if value is None:
         return Wikicode(SmartList())
     if hasattr(value, "read"):
-        return parse_anything(value.read(), context, skip_style_tags)
+        return parse_anything(value.read(), context, skip_style_tags=skip_style_tags)
     try:
         nodelist = SmartList()
         for item in value:
-            nodelist += parse_anything(item, context, skip_style_tags).nodes
+            nodelist += parse_anything(
+                item, context, skip_style_tags=skip_style_tags
+            ).nodes
         return Wikicode(nodelist)
     except TypeError as exc:
         error = (

tests/test_smart_list.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2012-2020 Ben Kurtovic <ben.kurtovic@gmail.com>
+# Copyright (C) 2012-2023 Ben Kurtovic <ben.kurtovic@gmail.com>
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -22,6 +22,8 @@
 Test cases for the SmartList class and its child, ListProxy.
 """
 
+import pickle
+
 import pytest
 
 from mwparserfromhell.smart_list import SmartList
@@ -432,3 +434,27 @@ def test_influence():
     assert [6, 5, 2, 3, 4, 1] == parent
     assert [4, 3, 2] == child2
     assert 0 == len(parent._children)
+
+
+@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1))
+def test_pickling(protocol: int):
+    """test SmartList objects behave properly when pickling"""
+    parent = SmartList([0, 1, 2, 3, 4, 5])
+    enc = pickle.dumps(parent, protocol=protocol)
+    assert pickle.loads(enc) == parent
+
+    child = parent[1:3]
+    assert len(parent._children) == 1
+    assert list(parent._children.values())[0][0]() is child
+    enc = pickle.dumps(parent, protocol=protocol)
+    parent2 = pickle.loads(enc)
+    assert parent2 == parent
+    assert parent2._children == {}
+
+    enc = pickle.dumps(child, protocol=protocol)
+    child2 = pickle.loads(enc)
+    assert child2 == child
+    assert child2._parent == parent
+    assert child2._parent is not parent
+    assert len(child2._parent._children) == 1
+    assert list(child2._parent._children.values())[0][0]() is child2

tests/test_template.py

@@ -725,7 +725,7 @@ def test_formatting():
         ),
     ]
 
-    for (original, expected) in tests:
+    for original, expected in tests:
         code = parse(original)
         template = code.filter_templates()[0]
         template.add("pop", "12345<ref>example ref</ref>")

tests/test_wikicode.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2012-2020 Ben Kurtovic <ben.kurtovic@gmail.com>
+# Copyright (C) 2012-2023 Ben Kurtovic <ben.kurtovic@gmail.com>
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -24,6 +24,7 @@ Tests for the Wikicode class, which manages a list of nodes.
 
 from functools import partial
 import re
+import pickle
 from types import GeneratorType
 
 import pytest
@@ -60,6 +61,14 @@ def test_nodes():
         code.__setattr__("nodes", object)
 
 
+@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1))
+def test_pickling(protocol: int):
+    """test Wikicode objects can be pickled"""
+    code = parse("Have a {{template}} and a [[page|link]]")
+    enc = pickle.dumps(code, protocol=protocol)
+    assert pickle.loads(enc) == code
+
+
 def test_get():
     """test Wikicode.get()"""
     code = parse("Have a {{template}} and a [[page|link]]")