1616
1717from  __future__ import  absolute_import , division , print_function 
1818
19+ from  io  import  TextIOWrapper 
1920from  collections  import  deque , namedtuple 
2021from  splunklib  import  six 
2122try :
3940
4041csv .field_size_limit (10485760 )  # The default value is 128KB; upping to 10MB. See SPL-12117 for background on this issue 
4142
42- # SPL-175233 -- python3 stdout is already binary 
43- if  sys .platform  ==  'win32'  and  sys .version_info  <=  (3 , 0 ):
44-     # Work around the fact that on Windows '\n' is mapped to '\r\n'. The typical solution is to simply open files in 
45-     # binary mode, but stdout is already open, thus this hack. 'CPython' and 'PyPy' work differently. We assume that 
46-     # all other Python implementations are compatible with 'CPython'. This might or might not be a valid assumption. 
47-     from  platform  import  python_implementation 
48-     implementation  =  python_implementation ()
49-     fileno  =  sys .stdout .fileno ()
50-     if  implementation  ==  'PyPy' :
51-         sys .stdout  =  os .fdopen (fileno , 'wb' , 0 )
52-     else :
53-         from  msvcrt  import  setmode 
54-         setmode (fileno , os .O_BINARY )
43+ 
44+ def  set_binary_mode (fh ):
45+     """ Helper method to set up binary mode for file handles. 
46+     Emphasis being sys.stdin, sys.stdout, sys.stderr. 
47+     For python3, we want to return .buffer 
48+     For python2+windows we want to set os.O_BINARY 
49+     """ 
50+     typefile  =  TextIOWrapper  if  sys .version_info  >=  (3 , 0 ) else  file 
51+     # check for file handle 
52+     if  not  isinstance (fh , typefile ):
53+         return  fh 
54+ 
55+     # check for python3 and buffer 
56+     if  sys .version_info  >=  (3 , 0 ) and  hasattr (fh , 'buffer' ):
57+         return  fh .buffer 
58+     # check for python3 
59+     elif  sys .version_info  >=  (3 , 0 ):
60+         pass 
61+     # check for windows python2. SPL-175233 -- python3 stdout is already binary 
62+     elif  sys .platform  ==  'win32' :
63+         # Work around the fact that on Windows '\n' is mapped to '\r\n'. The typical solution is to simply open files in 
64+         # binary mode, but stdout is already open, thus this hack. 'CPython' and 'PyPy' work differently. We assume that 
65+         # all other Python implementations are compatible with 'CPython'. This might or might not be a valid assumption. 
66+         from  platform  import  python_implementation 
67+         implementation  =  python_implementation ()
68+         if  implementation  ==  'PyPy' :
69+             return  os .fdopen (fh .fileno (), 'wb' , 0 )
70+         else :
71+             import  msvcrt 
72+             msvcrt .setmode (fh .fileno (), os .O_BINARY )
73+     return  fh 
5574
5675
5776class  CommandLineParser (object ):
@@ -349,6 +368,7 @@ class InputHeader(dict):
349368    """ Represents a Splunk input header as a collection of name/value pairs. 
350369
351370    """ 
371+ 
352372    def  __str__ (self ):
353373        return  '\n ' .join ([name  +  ':'  +  value  for  name , value  in  six .iteritems (self )])
354374
@@ -376,7 +396,8 @@ def read(self, ifile):
376396                # continuation of the current item 
377397                value  +=  urllib .parse .unquote (line )
378398
379-         if  name  is  not None : self [name ] =  value [:- 1 ] if  value [- 1 ] ==  '\n '  else  value 
399+         if  name  is  not None :
400+             self [name ] =  value [:- 1 ] if  value [- 1 ] ==  '\n '  else  value 
380401
381402
382403Message  =  namedtuple ('Message' , ('type' , 'text' ))
@@ -473,7 +494,7 @@ class RecordWriter(object):
473494    def  __init__ (self , ofile , maxresultrows = None ):
474495        self ._maxresultrows  =  50000  if  maxresultrows  is  None  else  maxresultrows 
475496
476-         self ._ofile  =  ofile 
497+         self ._ofile  =  set_binary_mode ( ofile ) 
477498        self ._fieldnames  =  None 
478499        self ._buffer  =  StringIO ()
479500
@@ -501,7 +522,13 @@ def ofile(self):
501522
502523    @ofile .setter  
503524    def  ofile (self , value ):
504-         self ._ofile  =  value 
525+         self ._ofile  =  set_binary_mode (value )
526+ 
527+     def  write (self , data ):
528+         bytes_type  =  bytes  if  sys .version_info  >=  (3 , 0 ) else  str 
529+         if  not  isinstance (data , bytes_type ):
530+             data  =  data .encode ('utf-8' )
531+         self .ofile .write (data )
505532
506533    def  flush (self , finished = None , partial = None ):
507534        assert  finished  is  None  or  isinstance (finished , bool )
@@ -661,19 +688,11 @@ class RecordWriterV1(RecordWriter):
661688
662689    def  flush (self , finished = None , partial = None ):
663690
664-         # SPL-175233 
665-         def  writeEOL ():
666-             if  sys .version_info  >=  (3 , 0 ) and  sys .platform  ==  'win32' :
667-                 write ('\n ' )
668-             else :
669-                 write ('\r \n ' )
670- 
671691        RecordWriter .flush (self , finished , partial )  # validates arguments and the state of this instance 
672692
673693        if  self ._record_count  >  0  or  (self ._chunk_count  ==  0  and  'messages'  in  self ._inspector ):
674694
675695            messages  =  self ._inspector .get ('messages' )
676-             write  =  self ._ofile .write 
677696
678697            if  self ._chunk_count  ==  0 :
679698
@@ -685,12 +704,12 @@ def writeEOL():
685704                    message_level  =  RecordWriterV1 ._message_level .get 
686705
687706                    for  level , text  in  messages :
688-                         write (message_level (level , level ))
689-                         write ('=' )
690-                         write (text )
691-                         writeEOL ( )
707+                         self . write (message_level (level , level ))
708+                         self . write ('=' )
709+                         self . write (text )
710+                         self . write ( ' \r \n ' 
692711
693-                 writeEOL ( )
712+                 self . write ( ' \r \n ' 
694713
695714            elif  messages  is  not None :
696715
@@ -708,7 +727,7 @@ def writeEOL():
708727                for  level , text  in  messages :
709728                    print (level , text , file = stderr )
710729
711-             write (self ._buffer .getvalue ())
730+             self . write (self ._buffer .getvalue ())
712731            self ._clear ()
713732            self ._chunk_count  +=  1 
714733            self ._total_record_count  +=  self ._record_count 
@@ -766,7 +785,7 @@ def write_metadata(self, configuration):
766785
767786        metadata  =  chain (six .iteritems (configuration ), (('inspector' , self ._inspector  if  self ._inspector  else  None ),))
768787        self ._write_chunk (metadata , '' )
769-         self ._ofile . write ('\n ' )
788+         self .write ('\n ' )
770789        self ._clear ()
771790
772791    def  write_metric (self , name , value ):
@@ -781,19 +800,22 @@ def _write_chunk(self, metadata, body):
781800
782801        if  metadata :
783802            metadata  =  str ('' .join (self ._iterencode_json (dict ([(n , v ) for  n , v  in  metadata  if  v  is  not None ]), 0 )))
803+             if  sys .version_info  >=  (3 , 0 ):
804+                 metadata  =  metadata .encode ('utf-8' )
784805            metadata_length  =  len (metadata )
785806        else :
786807            metadata_length  =  0 
787808
809+         if  sys .version_info  >=  (3 , 0 ):
810+             body  =  body .encode ('utf-8' )
788811        body_length  =  len (body )
789812
790813        if  not  (metadata_length  >  0  or  body_length  >  0 ):
791814            return 
792815
793816        start_line  =  'chunked 1.0,%s,%s\n '  %  (metadata_length , body_length )
794-         write  =  self ._ofile .write 
795-         write (start_line )
796-         write (metadata )
797-         write (body )
817+         self .write (start_line )
818+         self .write (metadata )
819+         self .write (body )
798820        self ._ofile .flush ()
799821        self ._flushed  =  False 
0 commit comments