2025-01-09T22:20:05Z

omlish/formats/xml.py

@@ -18,6 +18,49 @@
 ##
 
 
+def strip_ns(tag: str) -> str:
+    # It really do just be like this:
+    # https://github.com/python/cpython/blob/ff3bc82f7c9882c27aad597aac79355da7257186/Lib/xml/etree/ElementTree.py#L803-L804
+    if tag[:1] == '{':
+        _, tag = tag[1:].rsplit('}', 1)
+    return tag
+
+
+##
+
+
+ITER_PARSE_EVENTS = ('start', 'end', 'comment', 'pi', 'start-ns', 'end-ns')
+
+
+def yield_root_children(
+        source: ta.Any,
+        *,
+        retain_on_root: bool = False,
+        **kwargs: ta.Any,
+) -> ta.Iterator['ET.Element']:
+    it = iter(ET.iterparse(source, ('start', 'end'), **kwargs))
+
+    ev, root = next(it)
+    if ev != 'start':
+        raise RuntimeError(ev)
+    yield root
+
+    depth = 0
+    for ev, el in it:
+        if ev == 'start':
+            depth += 1
+
+        elif ev == 'end':
+            depth -= 1
+            if not depth:
+                if not retain_on_root:
+                    root.remove(el)
+                yield el
+
+
+##
+
+
 @dc.dataclass(frozen=True)
 class SimpleElement:
     tag: str

ommlx/wiki/utils/xml.py

@@ -14,6 +14,7 @@
 
 from omlish import check
 from omlish import lang
+from omlish.formats.xml import strip_ns
 
 
 if ta.TYPE_CHECKING:
@@ -30,17 +31,6 @@
 ##
 
 
-def strip_ns(tag: str) -> str:
-    # It really do just be like this:
-    # https://github.com/python/cpython/blob/ff3bc82f7c9882c27aad597aac79355da7257186/Lib/xml/etree/ElementTree.py#L803-L804
-    if tag[:1] == '{':
-        _, tag = tag[1:].rsplit('}', 1)
-    return tag
-
-
-##
-
-
 ITER_PARSE_EVENTS = ('start', 'end', 'comment', 'pi', 'start-ns', 'end-ns')