update modular inputs to use text TextIO consistently across 2.7 and 3.x

EventWriter will now wrap a non-TextIO object with a TextIO wrapper,
and ensures that all writing is done as string through six.ensure_text

Author: zenmoto

splunklib/modularinput/event.py

@@ -13,6 +13,9 @@
 # under the License.
 
 from __future__ import absolute_import
+from io import TextIOBase
+from splunklib.six import ensure_text
+
 try:
     import xml.etree.cElementTree as ET
 except ImportError as ie:
@@ -104,5 +107,8 @@ def write_to(self, stream):
         if self.done:
             ET.SubElement(event, "done")
 
-        stream.write(ET.tostring(event))
+        if isinstance(stream, TextIOBase):
+            stream.write(ensure_text(ET.tostring(event)))
+        else:
+            stream.write(ET.tostring(event))
         stream.flush()

splunklib/modularinput/event_writer.py

@@ -15,6 +15,8 @@
 from __future__ import absolute_import
 import sys
 
+from io import TextIOWrapper, TextIOBase
+from splunklib.six import ensure_text
 from .event import ET
 
 try:
@@ -24,7 +26,6 @@
 
 class EventWriter(object):
     """``EventWriter`` writes events and error messages to Splunk from a modular input.
-
     Its two important methods are ``writeEvent``, which takes an ``Event`` object,
     and ``log``, which takes a severity and an error message.
     """
@@ -42,8 +43,15 @@ def __init__(self, output = sys.stdout, error = sys.stderr):
         :param output: Where to write the output; defaults to sys.stdout.
         :param error: Where to write any errors; defaults to sys.stderr.
         """
-        self._out = output
-        self._err = error
+        if isinstance(output, TextIOBase):
+            self._out = output
+        else:
+            self._out = TextIOWrapper(output)
+
+        if isinstance(output, TextIOBase):
+            self._err = error
+        else:
+            self._err = TextIOWrapper(error)
 
         # has the opening <stream> tag been written yet?
         self.header_written = False
@@ -55,20 +63,18 @@ def write_event(self, event):
         """
 
         if not self.header_written:
-            self._out.write(b"<stream>")
+            self._out.write(ensure_text("<stream>"))
             self.header_written = True
 
         event.write_to(self._out)
 
     def log(self, severity, message):
         """Logs messages about the state of this modular input to Splunk.
         These messages will show up in Splunk's internal logs.
-
         :param severity: ``string``, severity of message, see severities defined as class constants.
         :param message: ``string``, message to log.
         """
-
-        self._err.write(("%s %s\n" % (severity, message)).encode('utf-8'))
+        self._err.write(ensure_text("%s %s\n" % (severity, message)))
         self._err.flush()
 
     def write_xml_document(self, document):
@@ -77,9 +83,11 @@ def write_xml_document(self, document):
 
         :param document: An ``ElementTree`` object.
         """
-        self._out.write(ET.tostring(document))
+        data = ET.tostring(document)
+        self._out.write(ensure_text(data))
         self._out.flush()
 
     def close(self):
         """Write the closing </stream> tag to make this XML well formed."""
-        self._out.write(b"</stream>")
+        self._out.write(ensure_text("</stream>"))
+        self._out.flush()

tests/modularinput/test_event.py

@@ -18,7 +18,7 @@
 from tests.modularinput.modularinput_testlib import unittest, xml_compare, data_open
 from splunklib.modularinput.event import Event, ET
 from splunklib.modularinput.event_writer import EventWriter
-from io import BytesIO
+from io import BytesIO, TextIOWrapper
 
 try:
     from splunklib.six.moves import cStringIO as StringIO
@@ -144,6 +144,9 @@ def test_write_xml_is_sane(self):
         out = BytesIO()
         err = BytesIO()
 
+        outwrap = TextIOWrapper(out)
+        errwrap = TextIOWrapper(err)
+
         ew = EventWriter(out, err)
 
         expected_xml = ET.parse(data_open("data/event_maximal.xml")).getroot()

tests/modularinput/test_script.py

@@ -21,12 +21,7 @@
 from splunklib.modularinput.script import Script
 from splunklib.modularinput.scheme import Scheme
 
-try:
-    from splunklib.six.moves import cStringIO as StringIO
-except ImportError:
-    from splunklib.six import StringIO
-
-from io import BytesIO
+from io import StringIO, BytesIO
 
 try:
     import xml.etree.cElementTree as ET