diff --git a/wikiextractor/cirrus-extract.py b/wikiextractor/cirrus-extract.py index 71ca6058..3428602c 100755 --- a/wikiextractor/cirrus-extract.py +++ b/wikiextractor/cirrus-extract.py @@ -42,6 +42,8 @@ import gzip import logging +from .WikiExtractor import NextFile, OutputSplitter + # Program version version = '3.0' @@ -49,72 +51,6 @@ # ---------------------------------------------------------------------- -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):