1# Copyright 2016 Google Inc. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14"""Support module for tests for yapf.""" 15 16import difflib 17import sys 18import unittest 19 20from yapf.pytree import blank_line_calculator 21from yapf.pytree import comment_splicer 22from yapf.pytree import continuation_splicer 23from yapf.pytree import pytree_unwrapper 24from yapf.pytree import pytree_utils 25from yapf.pytree import pytree_visitor 26from yapf.pytree import split_penalty 27from yapf.pytree import subtype_assigner 28from yapf.yapflib import identify_container 29from yapf.yapflib import style 30 31 32class YAPFTest(unittest.TestCase): 33 34 def __init__(self, *args): 35 super(YAPFTest, self).__init__(*args) 36 37 def assertCodeEqual(self, expected_code, code): 38 if code != expected_code: 39 msg = ['Code format mismatch:', 'Expected:'] 40 linelen = style.Get('COLUMN_LIMIT') 41 for line in expected_code.splitlines(): 42 if len(line) > linelen: 43 msg.append('!> %s' % line) 44 else: 45 msg.append(' > %s' % line) 46 msg.append('Actual:') 47 for line in code.splitlines(): 48 if len(line) > linelen: 49 msg.append('!> %s' % line) 50 else: 51 msg.append(' > %s' % line) 52 msg.append('Diff:') 53 msg.extend( 54 difflib.unified_diff( 55 code.splitlines(), 56 expected_code.splitlines(), 57 fromfile='actual', 58 tofile='expected', 59 lineterm='')) 60 self.fail('\n'.join(msg)) 61 62 63def ParseAndUnwrap(code, dumptree=False): 64 """Produces logical lines from the given code. 65 66 Parses the code into a tree, performs comment splicing and runs the 67 unwrapper. 68 69 Arguments: 70 code: code to parse as a string 71 dumptree: if True, the parsed pytree (after comment splicing) is dumped 72 to stderr. Useful for debugging. 73 74 Returns: 75 List of logical lines. 76 """ 77 tree = pytree_utils.ParseCodeToTree(code) 78 comment_splicer.SpliceComments(tree) 79 continuation_splicer.SpliceContinuations(tree) 80 subtype_assigner.AssignSubtypes(tree) 81 identify_container.IdentifyContainers(tree) 82 split_penalty.ComputeSplitPenalties(tree) 83 blank_line_calculator.CalculateBlankLines(tree) 84 85 if dumptree: 86 pytree_visitor.DumpPyTree(tree, target_stream=sys.stderr) 87 88 llines = pytree_unwrapper.UnwrapPyTree(tree) 89 for lline in llines: 90 lline.CalculateFormattingInformation() 91 92 return llines 93