diff --git a/tests/test_text_format.py b/tests/test_text_format.py
new file mode 100644
index 00000000..84342fa6
--- /dev/null
+++ b/tests/test_text_format.py
@@ -0,0 +1,163 @@
+"""
+Tests for the --text output format, --discard_empty flag, and the
+--json/--text mutual exclusivity added in the "Add options for a bare
+text format & removing empty documents" change.
+
+Background
+----------
+wikiextractor's default output wraps each article in a
+... XML-style header/footer.
+--json instead emits one JSON object per line. --text is a third,
+simpler option: just the body text (the same `text` list used by the
+other two formats), with no header/footer/metadata at all -- intended
+for corpus-building use cases (e.g. char-LM training data) where the
+wrapper is pure noise. --discard_empty skips writing anything at all
+for articles whose body extracts to nothing (redirects, category-only
+stubs, etc.), and applies uniformly across all three output formats
+(the check happens before the format branch, not inside it).
+
+Run with:
+ python -m unittest tests.test_text_format -v
+or, from the tests/ directory:
+ python -m unittest test_text_format -v
+"""
+
+import json
+import subprocess
+import sys
+import unittest
+from io import StringIO
+
+sys.path.insert(0, '..') # allow running directly from tests/ without installing
+
+from wikiextractor.extract import Extractor
+import wikiextractor.extract as ex
+
+
+class TextFormatTestCase(unittest.TestCase):
+ """Base class that resets wikiextractor's module-level template state
+ and output-format flags before each test, so tests don't leak state
+ into one another.
+ """
+
+ def setUp(self):
+ ex.templates.clear()
+ ex.templateCache.clear()
+ ex.redirects.clear()
+ ex.Extractor.templatePrefix = "Template:"
+ # Reset all three output-format flags to their defaults; each
+ # test sets only the ones it needs.
+ ex.Extractor.to_json = False
+ ex.Extractor.to_text = False
+ ex.Extractor.discard_empty = False
+
+ def make_extractor(self, article_text, article_id=1, title="Test Article"):
+ return Extractor(article_id, str(article_id),
+ f"https://test.wikipedia.org/wiki?curid={article_id}",
+ title, [article_text])
+
+ def extract_output(self, article_text, **flags):
+ for name, value in flags.items():
+ setattr(ex.Extractor, name, value)
+ extractor = self.make_extractor(article_text)
+ out = StringIO()
+ extractor.extract(out, html_safe=True)
+ return out.getvalue()
+
+
+class TextFormatOutputTests(TextFormatTestCase):
+ """--text should contain the article body and nothing else -- no
+ wrapper, no XML/JSON metadata.
+ """
+
+ def test_text_format_has_no_doc_tags(self):
+ article_text = "This is a plain paragraph of article text."
+ output = self.extract_output(article_text, to_text=True)
+
+ self.assertNotIn("", output)
+ self.assertIn("This is a plain paragraph of article text.", output)
+
+ def test_doc_format_still_has_tags_by_default(self):
+ # Regression check: adding the --text branch must not disturb the
+ # default (no flags set) format.
+ article_text = "This is a plain paragraph of article text."
+ output = self.extract_output(article_text) # all flags default False
+
+ self.assertIn("", output)
+ self.assertIn("This is a plain paragraph of article text.", output)
+
+ def test_json_format_unaffected_by_text_option(self):
+ # Regression check: --json must still produce valid JSON with no
+ # tags, unchanged by the new --text branch existing.
+ article_text = "This is a plain paragraph of article text."
+ output = self.extract_output(article_text, to_json=True)
+
+ self.assertNotIn(" tags and will have just the content instead. This is mutually
+exclusinve with --json
+
The program performs template expansion by preprocesssng the whole dump and
collecting template definitions.
"""
@@ -546,8 +550,11 @@ def main():
metavar="n[KMG]")
groupO.add_argument("-c", "--compress", action="store_true",
help="compress output files using bzip")
- groupO.add_argument("--json", action="store_true",
- help="write output in json format instead of the default format")
+ groupOFormat = groupO.add_mutually_exclusive_group()
+ groupOFormat.add_argument("--json", action="store_true",
+ help="write output in json format instead of the default format")
+ groupOFormat.add_argument("--text", action="store_true",
+ help="write output in text format (body only, no title) instead of the default format")
groupO.add_argument("--discard_empty", action="store_true",
help="discard empty articles (such as redirects) rather than writing just the title")
@@ -586,6 +593,7 @@ def main():
if args.html:
Extractor.keepLinks = True
Extractor.to_json = args.json
+ Extractor.to_text = args.text
Extractor.discard_empty = args.discard_empty
try:
@@ -610,6 +618,13 @@ def main():
if args.debug:
logger.setLevel(logging.DEBUG)
+ if args.json:
+ logger.debug("Outputting to json format")
+ elif args.text:
+ logger.debug("Outputting to text format")
+ else:
+ logger.debug("Outputting to format")
+
input_file = args.input
if not Extractor.keepLinks:
diff --git a/wikiextractor/extract.py b/wikiextractor/extract.py
index e23354f5..564f5f2d 100644
--- a/wikiextractor/extract.py
+++ b/wikiextractor/extract.py
@@ -923,6 +923,8 @@ class Extractor():
##
# Whether to produce json instead of the default output format.
to_json = False
+ # Whether to produce text instead of the default output format.
+ to_text = False
##
# Whether or not to discard empty (title only) documents
@@ -994,6 +996,9 @@ def extract(self, out, html_safe=True):
out_str = json.dumps(json_data)
out.write(out_str)
out.write('\n')
+ elif self.to_text:
+ out.write('\n'.join(text))
+ out.write('\n\n\n')
else:
header = '\n' % (self.id, self.url, self.title)
# Separate header from text with a newline.