-

cool-RR · cool-RR · commit 74fc6b356542 · 2014-01-06T00:27:40.000+02:00
diff --git a/source_py3/python_toolbox/file_tools.py b/source_py3/python_toolbox/file_tools.py
@@ -5,7 +5,6 @@
 
 import pathlib
 import re
-import contextlib
 
 from python_toolbox import context_management
 
@@ -49,23 +48,18 @@ def create_file_renaming_if_taken(path, mode='x',
                                   errors=None, newline=None):
     assert 'x' in mode
     current_path = pathlib.Path(path)
-    with contextlib.ExitStack() as exit_stack:
-        for i in range(N_MAX_ATTEMPTS):
-            try:
-                file = exit_stack.enter_context(
-                    current_path.open(mode, buffering=buffering,
-                                      encoding=encoding, errors=errors,
-                                      newline=newline)
-                )
-            except FileExistsError:
-                current_path = _get_next_name(current_path)
-            else:
-                return FileContainer(file)
-        else:
-            raise Exception("Exceeded {} tries, can't create file {}".format(
-                N_MAX_ATTEMPTS,
-                path
-            ))
+    for i in range(N_MAX_ATTEMPTS):
+        try:
+            return current_path.open(mode, buffering=buffering,
+                                     encoding=encoding, errors=errors,
+                                     newline=newline)
+        except FileExistsError:
+            current_path = _get_next_name(current_path)
+    else:
+        raise Exception("Exceeded {} tries, can't create file {}".format(
+            N_MAX_ATTEMPTS,
+            path
+        ))
     
 
 def write_to_file_renaming_if_taken(path, data, mode='x',
@@ -74,5 +68,6 @@ def write_to_file_renaming_if_taken(path, data, mode='x',
     with create_file_renaming_if_taken(
         path, mode=mode, buffering=buffering, encoding=encoding, errors=errors,
         newline=newline) as file:
-        file.write(data)
+        
+        return file.write(data)
         
diff --git a/source_py3/test_python_toolbox/test_file_tools/__init__.py b/source_py3/test_python_toolbox/test_file_tools/__init__.py
diff --git a/source_py3/test_python_toolbox/test_file_tools/test_renaming_file.py b/source_py3/test_python_toolbox/test_file_tools/test_renaming_file.py
@@ -0,0 +1,57 @@
+# Copyright 2009-2014 Ram Rachum.
+# This program is distributed under the MIT license.
+
+import pathlib
+
+import python_toolbox
+from python_toolbox import temp_file_tools
+
+from python_toolbox import file_tools
+
+
+def test():
+    with temp_file_tools.TemporaryFolder(prefix='test_python_toolbox_') as \
+                                                              temporary_folder:
+        get_file_names_set = \
+                  lambda: set(path.name for path in temporary_folder.glob('*'))
+        assert not get_file_names_set()
+        
+        file_path = temporary_folder / 'meow.txt'
+        string_to_write = "I'm a cat, hear me meow!"
+        
+        assert file_tools.write_to_file_renaming_if_taken(
+                          file_path, string_to_write) == len(string_to_write)
+        
+        with file_path.open('r') as file:
+            assert file.read() == string_to_write
+            
+        assert get_file_names_set() == {'meow.txt'}
+            
+        
+        assert file_tools.write_to_file_renaming_if_taken(
+                          file_path, string_to_write) == len(string_to_write)
+        assert file_tools.write_to_file_renaming_if_taken(
+                          file_path, string_to_write) == len(string_to_write)
+        assert file_tools.write_to_file_renaming_if_taken(
+                          file_path, string_to_write) == len(string_to_write)
+            
+        with (temporary_folder / 'meow (2).txt').open('r') as last_file_input:
+            assert last_file_input.read() == string_to_write
+        
+        assert get_file_names_set() == {'meow.txt', 'meow (1).txt',
+                                        'meow (2).txt', 'meow (3).txt'}
+        
+        with file_tools.create_file_renaming_if_taken(file_path) as last_file:
+            assert not last_file.closed
+            last_file.write(string_to_write[:5])
+            last_file.write(string_to_write[5:])
+            
+        assert last_file.closed
+        
+        assert get_file_names_set() == {'meow.txt', 'meow (1).txt',
+                                        'meow (2).txt', 'meow (3).txt',
+                                        'meow (4).txt'}
+        
+        with pathlib.Path(last_file.name).open('r') as last_file_input:
+            assert last_file_input.read() == string_to_write
+            
