Allowing input file to be same as output file.

remove.py

@@ -2,7 +2,7 @@
 
 def removeWatermark(input_fname: str, output_fname: str):
 
-    with open(input_fname, "rb") as input_file, open(output_fname, "wb") as output_file:
+    with open(input_fname, "rb") as input_file:
 
         reader = PdfReader(input_file)
         writer = PdfWriter()
@@ -14,14 +14,18 @@ def removeWatermark(input_fname: str, output_fname: str):
 
             writer.add_page(page)
 
+    with open(output_fname, "wb") as output_file:
+
         writer.write(output_file)
 
 if __name__ == "__main__":
 
     import sys
 
-    if len(sys.argv) != 3:
-        raise RuntimeError("Arguments not correct!")
-
-    removeWatermark(sys.argv[1], sys.argv[2])
+    if len(sys.argv) == 2:
+        removeWatermark(sys.argv[1], sys.argv[1])
+    elif len(sys.argv) == 3:
+        removeWatermark(sys.argv[1], sys.argv[2])
+    else:
+        raise RuntimeError("Wrong number of arguments!")