Windows file mode and color fixes for test.py, tooltest.py, gentest.py

mmgen/main.py

@@ -27,7 +27,10 @@ def launch(what):
 	if what == "keygen": what = "addrgen"
 
 	try: import termios
-	except: __import__("mmgen.main_" + what) # Windows
+	except: # Windows
+		from mmgen.util import start_mscolor
+		start_mscolor()
+		__import__("mmgen.main_" + what)
 	else:
 		import sys,atexit
 		fd = sys.stdin.fileno()

mmgen/test.py

@@ -22,17 +22,9 @@ test.py:  Shared routines for the test suites
 
 import sys,os
 from binascii import hexlify
-from mmgen.util import msg,write_to_file
+from mmgen.util import msg,write_to_file,red,green
 import mmgen.opt as opt
 
-_red,_grn,_yel,_cya,_reset = (
-	["\033[%sm" % c for c in "31;1","32;1","33;1","36;1","0"]
-)
-def red(s):    return _red+s+_reset
-def green(s):  return _grn+s+_reset
-def yellow(s): return _yel+s+_reset
-def cyan(s):   return _cya+s+_reset
-
 def cleandir(d):
 	try:    files = os.listdir(d)
 	except: return
@@ -57,8 +49,8 @@ def mk_tmpdir(cfg):
 def get_tmpfile_fn(cfg,fn):
 	return os.path.join(cfg['tmpdir'],fn)
 
-def write_to_tmpfile(cfg,fn,data):
-	write_to_file(os.path.join(cfg['tmpdir'],fn),data,silent=True)
+def write_to_tmpfile(cfg,fn,data,mode='wb'):
+	write_to_file(os.path.join(cfg['tmpdir'],fn),data,silent=True,mode=mode)
 
 def read_from_tmpfile(cfg,fn):
 	from mmgen.util import get_data_from_file

mmgen/tool.py

@@ -256,8 +256,11 @@ def hexdump(infile, cols=8, line_nums=True):
 			cols=cols, line_nums=line_nums))
 
 def unhexdump(infile):
+	if sys.platform[:3] == "win":
+		import msvcrt
+		msvcrt.setmode(sys.stdout.fileno(),os.O_BINARY)
 	sys.stdout.write(decode_pretty_hexdump(
-				get_data_from_file(infile,dash=True,silent=True)))
+			get_data_from_file(infile,dash=True,silent=True)))
 
 def strtob58(s):
 	enc = bitcoin.b58encode(s)
@@ -566,7 +569,7 @@ def rand2file(outfile, nbytes, threads=4, silent=False):
 	bsize = 2**20
 	roll = bsize * 4
 	if opt.outdir: outfile = make_full_path(opt.outdir,outfile)
-	f = open(outfile,"w")
+	f = open(outfile,"wb")
 
 	from Crypto.Cipher import AES
 	from Crypto.Util import Counter

mmgen/util.py

@@ -29,15 +29,27 @@ import mmgen.globalvars as g
 
 pnm = g.proj_name
 
-_red,_grn,_yel,_cya,_reset = (
-	["\033[%sm" % c for c in "31;1","32;1","33;1","36;1","0"],
-	["","","","",""]
-)[int(sys.platform[:3] == "win")]
-
-def red(s):    return _red+s+_reset
-def green(s):  return _grn+s+_reset
-def yellow(s): return _yel+s+_reset
-def cyan(s):   return _cya+s+_reset
+_red,_grn,_yel,_cya,_reset = \
+	["\033[%sm" % c for c in "31;1","32;1","33;1","36;1","0"]
+
+def red(s):     return _red+s+_reset
+def green(s):   return _grn+s+_reset
+def yellow(s):  return _yel+s+_reset
+def cyan(s):    return _cya+s+_reset
+def nocolor(s): return s
+
+def start_mscolor():
+	if sys.platform[:3] == "win":
+		global red,green,yellow,cyan,nocolor
+		import os
+		if "MMGEN_NOMSCOLOR" in os.environ:
+			red = green = yellow = cyan = nocolor
+		else:
+			try:
+				import colorama
+				colorama.init(strip=True,convert=True)
+			except:
+				red = green = yellow = cyan = nocolor
 
 def msg(s):    sys.stderr.write(s+"\n")
 def msg_r(s):  sys.stderr.write(s)
@@ -484,7 +496,8 @@ def write_to_file(
 		desc="data",
 		confirm_overwrite=False,
 		verbose=False,
-		silent=False
+		silent=False,
+		mode='wb'
 	):
 
 	if opt.outdir: outfile = make_full_path(opt.outdir,outfile)
@@ -498,7 +511,7 @@ def write_to_file(
 		else:
 			if not silent: msg("Overwriting file '%s'" % outfile)
 
-	f = open_file_or_exit(outfile,'wb')
+	f = open_file_or_exit(outfile,mode)
 	try:
 		f.write(data)
 	except:

test/gentest.py

@@ -12,9 +12,11 @@ from binascii import hexlify
 
 import mmgen.opt as opt
 import mmgen.globalvars as g
-from mmgen.util import msg,msg_r,mmsg,mdie,red,green,vmsg
+from mmgen.util import msg,msg_r,mmsg,mdie,red,green,vmsg,start_mscolor
 from mmgen.bitcoin import hextowif,privnum2addr
 
+start_mscolor()
+
 rounds = 100
 opts_data = {
 	'desc': "Test addresses generated by {} against output of 'keyconv'".format(g.proj_name),

test/test.py

@@ -17,16 +17,11 @@ scripts = (
 
 import mmgen.globalvars as g
 import mmgen.opt as opt
-from mmgen.util import mmsg,mdie,Msg,die,capfirst,write_data_to_file
-from mmgen.test import *
+from mmgen.util import *
 
-if sys.platform[:3] == "win":
-	try:
-		import colorama
-		colorama.init(strip=True,convert=True)
-	except:
-		def nocolor(s): return s
-		red = green = yellow = cyan = nocolor
+start_mscolor()
+
+from mmgen.test import *
 
 tb_cmd = "scripts/traceback.py"
 hincog_fn      = "rand_data"
@@ -420,8 +415,8 @@ cmd_args = opt.opts.init(opts_data)
 if opt.system: sys.path.pop(0)
 ni = bool(opt.non_interactive)
 
-# temporary
-#os.environ["MMGEN_USE_OLD_SCRIPTS"] = "1"
+# Disable MS color in spawned scripts due to bad interactions
+os.environ["MMGEN_NOMSCOLOR"] = "1"
 
 if opt.debug_scripts: os.environ["MMGEN_DEBUG"] = "1"
 

test/tooltest.py

@@ -9,9 +9,11 @@ os.chdir(os.path.join(pn,os.pardir))
 sys.path.__setitem__(0,os.path.abspath(os.curdir))
 
 import mmgen.opt as opt
-from mmgen.util import msg,msg_r,vmsg,vmsg_r,Msg,mmsg,mdie
+from mmgen.util import msg,msg_r,vmsg,vmsg_r,Msg,mmsg,mdie,start_mscolor
 from collections import OrderedDict
 
+start_mscolor()
+
 cmd_data = OrderedDict([
 	('util', {
 			'desc': "base conversion, hashing and file utilities",
@@ -150,12 +152,12 @@ class MMGenToolTestSuite(object):
 		self.__class__.__dict__[cmd](*([self,cmd] + file_list))
 
 
-	def run_cmd(self,name,tool_args,kwargs="",extra_msg="",silent=False):
+	def run_cmd(self,name,tool_args,kwargs="",extra_msg="",silent=False,strip=True):
 		mmgen_tool = "mmgen-tool"
 		if not opt.system:
 			mmgen_tool = os.path.join(os.curdir,mmgen_tool)
 
-		sys_cmd = [mmgen_tool, "-d",cfg['tmpdir'], name] + tool_args + kwargs.split()
+		sys_cmd = ["python", mmgen_tool, "-d",cfg['tmpdir'], name] + tool_args + kwargs.split()
 		if extra_msg: extra_msg = "(%s)" % extra_msg
 		full_name = " ".join([name]+kwargs.split()+extra_msg.split())
 		if not silent:
@@ -166,12 +168,13 @@ class MMGenToolTestSuite(object):
 				msg_r("Testing %-31s%s" % (full_name+":",""))
 
 		import subprocess
-		return subprocess.check_output(sys_cmd)
+		ret = subprocess.check_output(sys_cmd)
+		return (ret,ret.rstrip())[int(strip)]
 
 	def run_cmd_chk(self,name,f1,f2,kwargs="",extra_msg=""):
-		idata = read_from_file(f1)[:-1]
-		odata = read_from_file(f2)[:-1]
-		ret = self.run_cmd(name,[odata],kwargs=kwargs,extra_msg=extra_msg)[:-1]
+		idata = read_from_file(f1).rstrip()
+		odata = read_from_file(f2).rstrip()
+		ret = self.run_cmd(name,[odata],kwargs=kwargs,extra_msg=extra_msg)
 		vmsg("In:   " + repr(odata))
 		vmsg("Out:  " + repr(ret))
 		if ret == idata: ok()
@@ -182,32 +185,32 @@ class MMGenToolTestSuite(object):
 		return ret
 
 	def run_cmd_nochk(self,name,f1,kwargs=""):
-		odata = read_from_file(f1)[:-1]
-		ret = self.run_cmd(name,[odata],kwargs=kwargs)[:-1]
+		odata = read_from_file(f1).rstrip()
+		ret = self.run_cmd(name,[odata],kwargs=kwargs)
 		vmsg("In:   " + repr(odata))
 		vmsg("Out:  " + repr(ret))
 		return ret
 
 	def run_cmd_out(self,name,carg=None,Return=False,kwargs="",fn_idx="",extra_msg=""):
-		if carg: write_to_tmpfile(cfg,"%s%s.in" % (name,fn_idx),carg+"\n")
+		if carg: write_to_tmpfile(cfg,"%s%s.in" % (name,fn_idx),carg+"\n",mode='w')
 		ret = self.run_cmd(name,[carg] if carg else [],kwargs=kwargs,extra_msg=extra_msg)
 		if carg: vmsg("In:   " + repr(carg))
-		vmsg("Out:  " + repr(ret[:-1]))
+		vmsg("Out:  " + repr(ret))
 		if ret:
-			write_to_tmpfile(cfg,"%s%s.out" % (name,fn_idx),ret)
+			write_to_tmpfile(cfg,"%s%s.out" % (name,fn_idx),ret+"\n",mode='w')
 			if Return: return ret
 			else:   ok()
 		else:
 			msg(red("Error for command '%s'" % name))
 			sys.exit(3)
 
-	def run_cmd_randfileinput(self,name):
+	def run_cmd_randinput(self,name,strip=True):
 		s = os.urandom(128)
 		fn = name+".in"
 		write_to_tmpfile(cfg,fn,s)
-		ret = self.run_cmd(name,[get_tmpfile_fn(cfg,fn)])
+		ret = self.run_cmd(name,[get_tmpfile_fn(cfg,fn)],strip=strip)
 		fn = name+".out"
-		write_to_tmpfile(cfg,fn,ret)
+		write_to_tmpfile(cfg,fn,ret+"\n",mode='w')
 		ok()
 		vmsg("Returned: %s" % ret)
 
@@ -235,19 +238,19 @@ class MMGenToolTestSuite(object):
 		self.run_cmd(name,[])
 		ok()
 
-	def id6(self,name):     self.run_cmd_randfileinput(name)
-	def id8(self,name):     self.run_cmd_randfileinput(name)
-	def hexdump(self,name): self.run_cmd_randfileinput(name)
+	def id6(self,name):     self.run_cmd_randinput(name)
+	def id8(self,name):     self.run_cmd_randinput(name)
+	def hexdump(self,name): self.run_cmd_randinput(name,strip=False)
 
 	def unhexdump(self,name,fn1,fn2):
-		ret = self.run_cmd(name,[fn2])
+		ret = self.run_cmd(name,[fn2],strip=False)
 		orig = read_from_file(fn1)
 		cmp_or_die(orig,ret)
 
 	def rand2file(self,name):
 		of = name + ".out"
 		dlen = 1024
-		self.run_cmd(name,[of,str(1024),"threads=4","silent=1"])
+		self.run_cmd(name,[of,str(1024),"threads=4","silent=1"],strip=False)
 		d = read_from_tmpfile(cfg,of)
 		cmp_or_die(dlen,len(d))
 
@@ -262,14 +265,14 @@ class MMGenToolTestSuite(object):
 	def b32tohex(self,name,f1,f2): self.run_cmd_chk(name,f1,f2)
 	def b58randenc(self,name):
 		ret = self.run_cmd_out(name,Return=True)
-		ok_or_die(ret[:-1],is_b58_str,"base 58 string")
+		ok_or_die(ret,is_b58_str,"base 58 string")
 	def randhex(self,name):
 		ret = self.run_cmd_out(name,Return=True)
-		ok_or_die(ret[:-1],binascii.unhexlify,"hex string")
+		ok_or_die(ret,binascii.unhexlify,"hex string")
 	def randwif(self,name):
 		for n,k in enumerate(["","compressed=1"]):
 			ret = self.run_cmd_out(name,kwargs=k,Return=True,fn_idx=n+1)
-			ok_or_die(ret[:-1],is_wif,"WIF key")
+			ok_or_die(ret,is_wif,"WIF key")
 	def randpair(self,name):
 		for n,k in enumerate(["","compressed=1"]):
 			wif,addr = self.run_cmd_out(name,kwargs=k,Return=True,fn_idx=n+1).split()