Repository: google/yapf
Pull request #158 (2015-06-16 23:14:30)

Add in_place argument to FormatFile.
(For use as a module.)


Patch:
diff --git a/yapf/yapflib/yapf_api.py b/yapf/yapflib/yapf_api.py
index 8b141c4..3dd3e4b 100644
--- a/yapf/yapflib/yapf_api.py
+++ b/yapf/yapflib/yapf_api.py
@@ -42,6 +42,7 @@ from lib2to3.pgen2 import tokenize
 from yapf.yapflib import blank_line_calculator
 from yapf.yapflib import comment_splicer
 from yapf.yapflib import continuation_splicer
+from yapf.yapflib import file_resources
 from yapf.yapflib import py3compat
 from yapf.yapflib import pytree_unwrapper
 from yapf.yapflib import pytree_utils
@@ -55,28 +56,44 @@ def FormatFile(filename,
                style_config=None,
                lines=None,
                print_diff=False,
-               verify=True):
+               verify=True,
+               in_place=False):
   """Format a single Python file and return the formatted code.
 
   Arguments:
     filename: (unicode) The file to reformat.
     remaining arguments: see comment at the top of this module.
+    in_place: (bool) If True, write the reformatted code back to the file.
 
   Returns:
-    Pair of reformatted_code, encoding. reformatted_code is None if the file
-    doesn't exist.
+    Pair of (reformatted_code, encoding). reformatted_code is None if the file
+    is sucessfully written to (having used in_place). reformatted_code is a
+    diff if print_diff is True.
+
+  Raises:
+    IOError: raised if there was an error reading the file.
   """
   _CheckPythonVersion()
+
+  if in_place and print_diff:
+    raise ValueError("Cannot pass both in_place and print_diff.")
+
   original_source, encoding = ReadFile(filename, logging.warning)
-  if original_source is None:
-    return None, encoding
 
-  return FormatCode(original_source,
-                    style_config=style_config,
-                    filename=filename,
-                    lines=lines,
-                    print_diff=print_diff,
-                    verify=verify), encoding
+  reformatted_source = FormatCode(original_source,
+                                  style_config=style_config,
+                                  filename=filename,
+                                  lines=lines,
+                                  print_diff=print_diff,
+                                  verify=verify)
+  if in_place:
+    with py3compat.open_with_encoding(filename,
+                                      mode='w',
+                                      encoding=encoding) as fd:
+      fd.write(reformatted_source)
+      return None, encoding
+
+  return reformatted_source, encoding
 
 
 def FormatCode(unformatted_source,
@@ -154,7 +171,7 @@ def ReadFile(filename, logger=None):
     The contents of filename.
 
   Raises:
-    IOError: raised during an error if a logger is not specified.
+    IOError: raised if there was an error reading the file.
   """
   try:
     with open(filename, 'rb') as fd:


Test patch:
diff --git a/yapftests/yapf_test.py b/yapftests/yapf_test.py
index a99e330..0d58f09 100644
--- a/yapftests/yapf_test.py
+++ b/yapftests/yapf_test.py
@@ -186,6 +186,21 @@ class FormatFileTest(unittest.TestCase):
     diff = yapf_api.FormatFile(file1, print_diff=True)[0]
     self.assertTrue(u'+    pass' in diff)
 
+  def testFormatFileInPlace(self):
+    unformatted_code = "True==False\n"
+    formatted_code = "True == False\n"
+    file1 = self._MakeTempFileWithContents('testfile1.py', unformatted_code)
+    result, encoding = yapf_api.FormatFile(file1, in_place=True)
+    self.assertEqual(result, None)
+    with open(file1) as f:
+      self.assertCodeEqual(formatted_code, f.read())
+
+    self.assertRaises(ValueError, yapf_api.FormatFile, file1,
+                      in_place=True, print_diff=True)
+
+  def testNoFile(self):
+    self.assertRaises(IOError, yapf_api.FormatFile, "not_a_file.py")
+
 
 class CommandLineTest(unittest.TestCase):
   """Test how calling yapf from the command line acts."""