mediawiki_dump_tools/Mediawiki-Utilities/mw/xml_dump/iteration/iterator.py

   1 import io
   2
   3 from ...types import serializable
   4 from ...util import none_or
   5 from ..element_iterator import ElementIterator
   6 from ..errors import MalformedXML
   7 from .namespace import Namespace
   8 from .page import Page
   9
  10
  11 class ConcatinatingTextReader(io.TextIOBase):
  12
  13     def __init__(self, *items):
  14         self.items = [io.StringIO(i) if isinstance(i, str) else i
  15                       for i in items]
  16
  17     def read(self, size=-1):
  18         return "".join(self._read(size))
  19
  20     def readline(self):
  21
  22         if len(self.items) > 0:
  23             line = self.items[0].readline()
  24             if line == "": self.items.pop(0)
  25         else:
  26             line = ""
  27
  28         return line
  29
  30     def _read(self, size):
  31         if size > 0:
  32             while len(self.items) > 0:
  33                 byte_vals = self.items[0].read(size)
  34                 yield byte_vals
  35                 if len(byte_vals) < size:
  36                     size = size - len(byte_vals) # Decrement bytes
  37                     self.items.pop(0)
  38                 else:
  39                     break
  40
  41         else:
  42             for item in self.items:
  43                 yield item.read()
  44
  45
  46
  47
  48 def concat(*stream_items):
  49     return ConcatinatingTextReader(*stream_items)
  50
  51
  52 class Iterator(serializable.Type):
  53     """
  54     XML Dump Iterator. Dump file meta data and a
  55     :class:`~mw.xml_dump.Page` iterator.  Instances of this class can be
  56     called as an iterator directly.  E.g.::
  57
  58         from mw.xml_dump import Iterator
  59
  60         # Construct dump file iterator
  61         dump = Iterator.from_file(open("example/dump.xml"))
  62
  63         # Iterate through pages
  64         for page in dump:
  65
  66             # Iterate through a page's revisions
  67             for revision in page:
  68
  69                 print(revision.id)
  70
  71     """
  72     __slots__ = ('site_name', 'base', 'generator', 'case', 'namespaces',
  73                  '__pages')
  74
  75     def __init__(self, site_name=None, dbname=None, base=None, generator=None,
  76                  case=None, namespaces=None, pages=None):
  77
  78         self.site_name = none_or(site_name, str)
  79         """
  80         The name of the site. : str | `None` (if not specified in the XML)
  81         """
  82
  83         self.dbname = none_or(dbname, str)
  84         """
  85         The database name of the site. : str | `None` (if not specified in the
  86         XML)
  87         """
  88
  89         self.base = none_or(base, str)
  90         """
  91         TODO: ??? : str | `None` (if not specified in the XML)
  92         """
  93
  94         self.generator = none_or(generator, str)
  95         """
  96         TODO: ??? : str | `None` (if not specified in the XML)
  97         """
  98
  99         self.case = none_or(case, str)
 100         """
 101         TODO: ??? : str | `None` (if not specified in the XML)
 102         """
 103
 104         self.namespaces = none_or(namespaces, list)
 105         """
 106         A list of :class:`mw.Namespace` | `None` (if not specified in the XML)
 107         """
 108
 109         # Should be a lazy generator of page info
 110         self.__pages = pages
 111
 112     def __iter__(self):
 113         return self.__pages
 114
 115     def __next__(self):
 116         return next(self.__pages)
 117
 118     @classmethod
 119     def load_namespaces(cls, element):
 120         namespaces = []
 121         for sub_element in element:
 122             tag = sub_element.tag
 123
 124             if tag == "namespace":
 125                 namespace = Namespace.from_element(sub_element)
 126                 namespaces.append(namespace)
 127             else:
 128                 assert False, "This should never happen"
 129
 130         return namespaces
 131
 132     @classmethod
 133     def load_site_info(cls, element):
 134
 135         site_name = None
 136         dbname = None
 137         base = None
 138         generator = None
 139         case = None
 140         namespaces = {}
 141
 142         for sub_element in element:
 143             if sub_element.tag == 'sitename':
 144                 site_name = sub_element.text
 145             if sub_element.tag == 'dbname':
 146                 dbname = sub_element.text
 147             elif sub_element.tag == 'base':
 148                 base = sub_element.text
 149             elif sub_element.tag == 'generator':
 150                 generator = sub_element.text
 151             elif sub_element.tag == 'case':
 152                 case = sub_element.text
 153             elif sub_element.tag == 'namespaces':
 154                 namespaces = cls.load_namespaces(sub_element)
 155
 156         return site_name, dbname, base, generator, case, namespaces
 157
 158     @classmethod
 159     def load_pages(cls, element):
 160
 161         for sub_element in element:
 162             tag = sub_element.tag
 163
 164             if tag == "page":
 165                 yield Page.from_element(sub_element)
 166             else:
 167                 assert MalformedXML("Expected to see 'page'.  " +
 168                                     "Instead saw '{0}'".format(tag))
 169
 170     @classmethod
 171     def from_element(cls, element):
 172
 173         site_name = None
 174         base = None
 175         generator = None
 176         case = None
 177         namespaces = None
 178
 179         # Consume <siteinfo>
 180         for sub_element in element:
 181             tag = sub_element.tag
 182             if tag == "siteinfo":
 183                 site_name, dbname, base, generator, case, namespaces = \
 184                     cls.load_site_info(sub_element)
 185                 break
 186
 187         # Consume all <page>
 188         pages = cls.load_pages(element)
 189
 190         return cls(site_name, dbname, base, generator, case, namespaces, pages)
 191
 192     @classmethod
 193     def from_file(cls, f):
 194         element = ElementIterator.from_file(f)
 195         assert element.tag == "mediawiki"
 196         return cls.from_element(element)
 197
 198     @classmethod
 199     def from_string(cls, string):
 200         f = io.StringIO(string)
 201         element = ElementIterator.from_file(f)
 202         assert element.tag == "mediawiki"
 203         return cls.from_element(element)
 204
 205     @classmethod
 206     def from_page_xml(cls, page_xml):
 207         header = """
 208         <mediawiki xmlns="http://www.mediawiki.org/xml/export-0.5/"
 209                    xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
 210                    xsi:schemaLocation="http://www.mediawiki.org/xml/export-0.5/
 211                      http://www.mediawiki.org/xml/export-0.5.xsd" version="0.5"
 212                    xml:lang="en">
 213         <siteinfo>
 214             <namespaces>
 215             </namespaces>
 216         </siteinfo>
 217         """
 218
 219         footer = "</mediawiki>"
 220
 221         return cls.from_file(concat(header, page_xml, footer))