Import NextFile and OutputSplitter from WikiExtractor

wikiextractor/cirrus-extract.py

@@ -42,79 +42,15 @@
 import gzip
 import logging
 
+from .WikiExtractor import NextFile, OutputSplitter
+
 # Program version
 version = '3.0'
 
 urlbase = 'http://it.wikipedia.org/'
 
 # ----------------------------------------------------------------------
 
-class NextFile(object):
-    """
-    Synchronous generation of next available file name.
-    """
-
-    filesPerDir = 100
-
-    def __init__(self, path_name):
-        self.path_name = path_name
-        self.dir_index = -1
-        self.file_index = -1
-
-    def next(self):
-        self.file_index = (self.file_index + 1) % NextFile.filesPerDir
-        if self.file_index == 0:
-            self.dir_index += 1
-        dirname = self._dirname()
-        if not os.path.isdir(dirname):
-            os.makedirs(dirname)
-        return self._filepath()
-
-    def _dirname(self):
-        char1 = self.dir_index % 26
-        char2 = int(self.dir_index / 26) % 26
-        return os.path.join(self.path_name, '%c%c' % (ord('A') + char2, ord('A') + char1))
-
-    def _filepath(self):
-        return '%s/wiki_%02d' % (self._dirname(), self.file_index)
-
-class OutputSplitter(object):
-    """
-    File-like object, that splits output to multiple files of a given max size.
-    """
-
-    def __init__(self, nextFile, max_file_size=0, compress=True):
-        """
-        :param nextfile: a NextFile object from which to obtain filenames
-            to use.
-        :param max_file_size: the maximum size of each file.
-        :para compress: whether to write data with bzip compression.
-        """
-        self.nextFile = nextFile
-        self.compress = compress
-        self.max_file_size = max_file_size
-        self.file = self.open(self.nextFile.next())
-
-    def reserve(self, size):
-        if self.file.tell() + size > self.max_file_size:
-            self.close()
-            self.file = self.open(self.nextFile.next())
-
-    def write(self, data):
-        self.reserve(len(data))
-        self.file.write(data)
-
-    def close(self):
-        self.file.close()
-
-    def open(self, filename):
-        if self.compress:
-            return bz2.BZ2File(filename + '.bz2', 'w')
-        else:
-            return open(filename, 'w')
-
-# ----------------------------------------------------------------------
-
 class Extractor(object):
 
     def extract(self, out):