handle multi-step redirects correctly

ikuyamada · ikuyamada · commit f4c6fda6b6f2 · 2024-01-02T21:07:19.000+09:00
diff --git a/wikipedia2vec/dump_db.py b/wikipedia2vec/dump_db.py
@@ -104,13 +104,24 @@ def redirects(self) -> Iterator[Tuple[str, str]]:
             for key, value in iter(cur):
                 yield (key.decode("utf-8"), value.decode("utf-8"))
 
-    def resolve_redirect(self, title: str) -> str:
-        with self._env.begin(db=self._redirect_db) as txn:
-            value = txn.get(title.encode("utf-8"))
-            if value:
-                return value.decode("utf-8")
-            else:
-                return title
+    def resolve_redirect(self, title: str, max_steps: int = 10) -> str:
+        visited = set([title])
+        cur_title = title
+        for _ in range(max_steps):
+            with self._env.begin(db=self._redirect_db) as txn:
+                value = txn.get(cur_title.encode("utf-8"))
+                if value:
+                    cur_title = value.decode("utf-8")
+                    if cur_title in visited:
+                        logger.warn(f"Detected redirect loop: {title}")
+                        return title
+                    visited.add(cur_title)
+
+                else:
+                    return cur_title
+
+        logger.warn(f"Max steps ({max_steps}) exceeded when resolving redirect: {title}")
+        return cur_title
 
     def is_redirect(self, title: str) -> bool:
         with self._env.begin(db=self._redirect_db) as txn:
