From e70c4c295f84f2fccffebccfd9e99a4a249c8fe4 Mon Sep 17 00:00:00 2001
From: George Litvinov <geolit@google.com>
Date: Tue, 25 Aug 2026 10:52:45 -0700
Subject: [PATCH] Implement picklemagic error handling. Add pydoc functions to
 unsafe strings

PiperOrigin-RevId: 970659223
---
 lib/constants.py                   |  2 ++
 third_party/corrupy/picklemagic.py | 40 +++++++++++++++++++++---------
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/lib/constants.py b/lib/constants.py
index e4bffe8..2148ff1 100644
--- a/lib/constants.py
+++ b/lib/constants.py
@@ -129,7 +129,9 @@
     "pty",
     "numpy.lib.npyio.loadtxt",
     "pycrypto",
+    "pydoc.locate",
     "pydoc.pipepager",
+    "pydoc.replace",
     "pandas.read_pickle",
     "python",
     "pywin32_system32",
diff --git a/third_party/corrupy/picklemagic.py b/third_party/corrupy/picklemagic.py
index 75119c3..d1e95b6 100644
--- a/third_party/corrupy/picklemagic.py
+++ b/third_party/corrupy/picklemagic.py
@@ -841,29 +841,45 @@ def load_build(self, *unused_args):
       return
 
     try:
-      inst.__setstate__(state)
-    except AttributeError:
+      setstate = getattr(inst, "__setstate__", None)
+      if setstate is not None:
+        setstate(state)
+        return
+
       logger.warning(
           "Attribute __setstate__ is not available for object %s", type(inst)
       )
       # Standard pickle fallback: if __setstate__ is not defined,
       # update __dict__ or slots.
+      dict_state = state
+      slots_state = None
       if isinstance(state, tuple) and len(state) == 2:
         dict_state, slots_state = state
-        if isinstance(dict_state, dict):
-          inst.__dict__.update(dict_state)
-        if isinstance(slots_state, dict):
-          for slot, val in slots_state.items():
-            setattr(inst, slot, val)
-      elif isinstance(state, dict):
-        logger.info("Updating __dict__ of %s with state", type(inst))
-        inst.__dict__.update(state)
-      else:
+
+      if isinstance(dict_state, dict):
+        inst_dict = getattr(inst, "__dict__", None)
+        if isinstance(inst_dict, dict):
+          logger.info("Updating __dict__ of %s with state", type(inst))
+          inst_dict.update(dict_state)
+        elif inst_dict is not None:
+          for k, v in dict_state.items():
+            try:
+              setattr(inst, k, v)
+            except Exception:
+              pass
+      elif dict_state is not None:
         logger.warning(
             "Cannot update state of %s with state of type %s",
             type(inst),
-            type(state),
+            type(dict_state),
         )
+
+      if isinstance(slots_state, dict):
+        for slot, val in slots_state.items():
+          try:
+            setattr(inst, slot, val)
+          except Exception:
+            pass
     except Exception as e:  # pylint: disable=broad-exception-caught
       logger.warning("Failed to set state on %s: %s", type(inst), e)
       logger.info("Proceeding after state failure on %s", type(inst))