From 25d6729e03b2483eddc29f46c40efd7b446ea3b3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?R=C3=A9mi=20BERTHO?= <remi.bertho@dalan.fr>
Date: Mon, 16 Oct 2017 22:51:03 +0200
Subject: [PATCH] Some improvments
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Rémi BERTHO <remi.bertho@dalan.fr>
---
 SSnR.py | 34 +++++++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/SSnR.py b/SSnR.py
index 5d47f65..408a84e 100755
--- a/SSnR.py
+++ b/SSnR.py
@@ -33,23 +33,24 @@ def main():
 	parser = argparse.ArgumentParser(description='Search and replace tool', prog='SSnR')
 	parser.add_argument('-e', '--regex', help='Regex', required=True)
 	parser.add_argument('-s', '--string', help='String', required=False)
-	parser.add_argument('-i', '--input', help='Input file', required=False, type=open)
-	parser.add_argument('-o', '--output', help='Output file', required=False,
-	                    type=argparse.FileType('w'))
+	parser.add_argument('-i', '--input', help='Input file', required=False)
+	parser.add_argument('-o', '--output', help='Output file', required=False)
 	parser.add_argument('-r', '--replace', help='Replace', required=False)
 	parser.add_argument('-m', '--print_nb_match', help='Print the number of match in replace',
 	                    required=False, action='store_true')
 	parser.add_argument('-c', '--ignore_case', help='Ignore the case',
 	                    required=False, action='store_true')
-	try:
-		args = vars(parser.parse_args())
-	except FileNotFoundError as exception:
-		print("Error: file not found: " + str(exception))
-		return -1
+
+	args = vars(parser.parse_args())
 
 	# Get input
 	if args["input"] is not None:
-		string = args["input"].read()
+		try:
+			with open(args["input"], "r") as file:
+				string = file.read()
+		except OSError as exception:
+			print("Error: file not found: " + str(exception))
+			return -1
 		is_file = True
 	elif args["string"] is not None:
 		string = args["string"]
@@ -58,6 +59,17 @@ def main():
 		print("Error: You need an input string or file")
 		return -1
 
+	# Get output
+	if args["output"] is not None:
+		try:
+			output_file = open(args["output"], "w")
+		except OSError as exception:
+			print("Error: file not found: " + str(exception))
+			return -1
+		use_output_file = True
+	else:
+		use_output_file = False
+
 	# Compile regex
 	try:
 		ex = compile_regex(args["regex"])
@@ -71,8 +83,8 @@ def main():
 	# Search or replace
 	if args["replace"] is not None:
 		replace_string = replace(ex, string, args["replace"], args["print_nb_match"])
-		if args["output"] is not None:
-			args["output"].write(replace_string)
+		if use_output_file:
+			output_file.write(replace_string)
 		else:
 			print(replace_string)
 	else: