Merge pull request pywinauto#1049 from eltimen/refactor-dump-tree

Rewrite dump_tree method

pywinauto/base_application.py

@@ -82,6 +82,7 @@
 import time
 import locale
 import codecs
+import collections
 
 import six
 
@@ -576,64 +577,117 @@ def _ctrl_identifiers(self):
 
         return control_name_map
 
-    def print_control_identifiers(self, depth=None, filename=None):
+    def dump_tree(self, depth=10, max_width=10, filename=None):
         """
-        Prints the 'identifiers'
+        Dump the 'identifiers' to console or a file
 
-        Prints identifiers for the control and for its descendants to
+        Dump identifiers for the control and for its descendants to
         a depth of **depth** (the whole subtree if **None**).
 
-        .. note:: The identifiers printed by this method have been made
+        :param depth: Max depth level of an element tree to dump (None: unlimited).
+
+        :param max_width: Max number of children of each element to dump (None: unlimited).
+
+        :param filename: Save tree to a specified file (None: print to stdout).
+
+        .. note:: The identifiers dumped by this method have been made
                unique. So if you have 2 edit boxes, they won't both have "Edit"
                listed in their identifiers. In fact the first one can be
                referred to as "Edit", "Edit0", "Edit1" and the 2nd should be
                referred to as "Edit2".
         """
         if depth is None:
             depth = sys.maxsize
+        if max_width is None:
+            max_width = sys.maxsize
         # Wrap this control
         this_ctrl = self.__resolve_control(self.criteria)[-1]
 
-        # Create a list of this control and all its descendants
-        all_ctrls = [this_ctrl, ] + this_ctrl.descendants(depth=depth)
-
-        # Create a list of all visible text controls
-        txt_ctrls = [ctrl for ctrl in all_ctrls if ctrl.can_be_label and ctrl.is_visible() and ctrl.window_text()]
-
-        # Build a dictionary of disambiguated list of control names
-        name_ctrl_id_map = findbestmatch.UniqueDict()
-        for index, ctrl in enumerate(all_ctrls):
-            ctrl_names = findbestmatch.get_control_names(ctrl, all_ctrls, txt_ctrls)
-            for name in ctrl_names:
-                name_ctrl_id_map[name] = index
-
-        # Swap it around so that we are mapped off the control indices
-        ctrl_id_name_map = {}
-        for name, index in name_ctrl_id_map.items():
-            ctrl_id_name_map.setdefault(index, []).append(name)
-
-        def print_identifiers(ctrls, current_depth=0, log_func=print):
+        ElementTreeNode = collections.namedtuple('ElementTreeNode', ['elem', 'id', 'children'])
+
+        def create_element_tree(element_list):
+            """Build elements tree and create list with pre-order tree traversal"""
+            depth_limit_reached = False
+            width_limit_reached = False
+            current_id = 0
+            elem_stack = collections.deque([(this_ctrl, None, 0)])
+            root_node = ElementTreeNode(this_ctrl, current_id, [])
+            while elem_stack:
+                current_elem, current_elem_parent_children, current_node_depth = elem_stack.pop()
+                if current_elem is None:
+                    elem_node = ElementTreeNode(None, current_id, [])
+                    current_elem_parent_children.append(elem_node)
+                else:
+                    if current_node_depth <= depth:
+                        if current_elem_parent_children is not None:
+                            current_id += 1
+                            elem_node = ElementTreeNode(current_elem, current_id, [])
+                            current_elem_parent_children.append(elem_node)
+                            element_list.append(current_elem)
+                        else:
+                            elem_node = root_node
+                        child_elements = current_elem.children()
+                        if len(child_elements) > max_width and current_node_depth < depth:
+                            elem_stack.append((None, elem_node.children, current_node_depth + 1))
+                            width_limit_reached = True
+                        for i in range(min(len(child_elements) - 1, max_width - 1), -1, -1):
+                            elem_stack.append((child_elements[i], elem_node.children, current_node_depth + 1))
+                    else:
+                        depth_limit_reached = True
+            return root_node, depth_limit_reached, width_limit_reached
+
+        # Create a list of this control, all its descendants
+        all_ctrls = [this_ctrl]
+
+        # Build element tree
+        elements_tree, depth_limit_reached, width_limit_reached = create_element_tree(all_ctrls)
+
+        show_best_match_names = self.allow_magic_lookup and not (depth_limit_reached or width_limit_reached)
+        if show_best_match_names:
+            # Create a list of all visible text controls
+            txt_ctrls = [ctrl for ctrl in all_ctrls if ctrl.can_be_label and ctrl.is_visible() and ctrl.window_text()]
+
+            # Build a dictionary of disambiguated list of control names
+            name_ctrl_id_map = findbestmatch.UniqueDict()
+            for index, ctrl in enumerate(all_ctrls):
+                ctrl_names = findbestmatch.get_control_names(ctrl, all_ctrls, txt_ctrls)
+                for name in ctrl_names:
+                    name_ctrl_id_map[name] = index
+
+            # Swap it around so that we are mapped off the control indices
+            ctrl_id_name_map = {}
+            for name, index in name_ctrl_id_map.items():
+                ctrl_id_name_map.setdefault(index, []).append(name)
+
+        def print_identifiers(element_node, current_depth=0, log_func=print):
             """Recursively print ids for ctrls and their descendants in a tree-like format"""
-            if len(ctrls) == 0 or current_depth > depth:
-                return
+            if current_depth == 0:
+                if depth_limit_reached:
+                    log_func('Warning: the whole hierarchy does not fit into depth={}. '
+                             'Increase depth parameter value or set it to None (unlimited, '
+                             'may freeze in case of very large number of elements).'.format(depth))
+                if self.allow_magic_lookup and not show_best_match_names:
+                    log_func('If the whole hierarchy fits into depth and max_width values, '
+                             'best_match names are dumped.')
+                log_func("Control Identifiers:")
 
             indent = current_depth * u"   | "
-            for ctrl in ctrls:
-                try:
-                    ctrl_id = all_ctrls.index(ctrl)
-                except ValueError:
-                    continue
+            output = indent + u'\n'
+
+            ctrl = element_node.elem
+            if ctrl is not None:
+                ctrl_id = element_node.id
                 ctrl_text = ctrl.window_text()
                 if ctrl_text:
                     # transform multi-line text to one liner
                     ctrl_text = ctrl_text.replace('\n', r'\n').replace('\r', r'\r')
+                output += indent + u"{class_name} - '{text}'    {rect}" \
+                                   "".format(class_name=ctrl.friendly_class_name(),
+                                             text=ctrl_text,
+                                             rect=ctrl.rectangle())
 
-                output = indent + u'\n'
-                output += indent + u"{class_name} - '{text}'    {rect}\n"\
-                    "".format(class_name=ctrl.friendly_class_name(),
-                              text=ctrl_text,
-                              rect=ctrl.rectangle())
-                output += indent + u'{}'.format(ctrl_id_name_map[ctrl_id])
+                if show_best_match_names:
+                    output += u'\n' + indent + u'{}'.format(ctrl_id_name_map[ctrl_id])
 
                 class_name = ctrl.class_name()
                 auto_id = None
@@ -653,28 +707,32 @@ def print_identifiers(ctrls, current_depth=0, log_func=print):
                     criteria_texts.append(u'control_type="{}"'.format(control_type))
                 if ctrl_text or class_name or auto_id:
                     output += u'\n' + indent + u'child_window(' + u', '.join(criteria_texts) + u')'
+            else:
+                output += indent + u'**********\n'
+                output += indent + u'Max children output limit ({}) has been reached. ' \
+                                   u'Set a larger max_width value or use max_width=None ' \
+                                   u'to see all children.\n'.format(max_width)
+                output += indent + u'**********'
+
+            if six.PY3:
+                log_func(output)
+            else:
+                log_func(output.encode(locale.getpreferredencoding(), errors='backslashreplace'))
 
-                if six.PY3:
-                    log_func(output)
-                else:
-                    log_func(output.encode(locale.getpreferredencoding(), errors='backslashreplace'))
-
-                print_identifiers(ctrl.children(), current_depth + 1, log_func)
+            if current_depth <= depth:
+                for child_elem in element_node.children:
+                    print_identifiers(child_elem, current_depth + 1, log_func)
 
         if filename is None:
-            print("Control Identifiers:")
-            print_identifiers([this_ctrl, ])
+            print_identifiers(elements_tree)
         else:
-            log_file = codecs.open(filename, "w", locale.getpreferredencoding())
-
-            def log_func(msg):
-                log_file.write(str(msg) + os.linesep)
-            log_func("Control Identifiers:")
-            print_identifiers([this_ctrl, ], log_func=log_func)
-            log_file.close()
+            with codecs.open(filename, "w", locale.getpreferredencoding()) as log_file:
+                def log_func(msg):
+                    log_file.write(str(msg) + os.linesep)
+                print_identifiers(elements_tree, log_func=log_func)
 
-    print_ctrl_ids = print_control_identifiers
-    dump_tree = print_control_identifiers
+    print_control_identifiers = deprecated(dump_tree, deprecated_name='print_control_identifiers')
+    print_ctrl_ids = deprecated(dump_tree, deprecated_name='print_ctrl_ids')
 
 
 #=========================================================================

pywinauto/unittests/test_application.py

@@ -1169,15 +1169,15 @@ def test_depth(self):
             len(self.app['Font'].descendants(depth=1)),
             len(self.app['Font'].descendants(depth=2)))
 
-    def test_print_control_identifiers(self):
-        """Make sure print_control_identifiers() doesn't crash"""
-        self.dlgspec.print_control_identifiers()
-        self.ctrlspec.print_control_identifiers()
-
-    def test_print_control_identifiers_file_output(self):
-        """Make sure print_control_identifiers() creates correct file"""
-        output_filename = "test_print_control_identifiers.txt"
-        self.dlgspec.print_ctrl_ids(filename=output_filename)
+    def test_dump_tree(self):
+        """Make sure dump_tree() doesn't crash"""
+        self.dlgspec.dump_tree()
+        self.ctrlspec.dump_tree()
+
+    def test_dump_tree_file_output(self):
+        """Make sure dump_tree() creates correct file"""
+        output_filename = "test_dump_tree.txt"
+        self.dlgspec.dump_tree(filename=output_filename)
         if os.path.isfile(output_filename):
             with open(output_filename, "r") as test_log_file:
                 content = str(test_log_file.readlines())
@@ -1186,7 +1186,7 @@ def test_print_control_identifiers_file_output(self):
                 self.assertTrue("child_window(class_name=\"msctls_statusbar32\"" in content)
             os.remove(output_filename)
         else:
-            self.fail("print_control_identifiers can't create a file")
+            self.fail("dump_tree can't create a file")
 
         self.ctrlspec.dump_tree(filename=output_filename)
         if os.path.isfile(output_filename):
@@ -1195,7 +1195,7 @@ def test_print_control_identifiers_file_output(self):
                 self.assertTrue("child_window(class_name=\"Edit\")" in content)
             os.remove(output_filename)
         else:
-            self.fail("print_control_identifiers can't create a file")
+            self.fail("dump_tree can't create a file")
 
     def test_find_elements_re(self):
         """Test for bug #90: A crash in 'find_elements' when called with 'title_re' argument"""
@@ -1324,7 +1324,7 @@ def test_folder_list(self):
             files_list = self.desktop.MFC_samplesDialog.Shell_Folder_View.Items_View.find()
             self.assertEqual([item.window_text() for item in files_list.get_items()],
                              [u'x64', u'BCDialogMenu.exe', u'CmnCtrl1.exe', u'CmnCtrl2.exe', u'CmnCtrl3.exe',
-                              u'CtrlTest.exe', u'mfc100u.dll', u'RebarTest.exe', u'RowList.exe', u'TrayMenu.exe'])
+                              u'CtrlTest.exe', u'mfc100u.dll', u'NewControls.exe', u'RebarTest.exe', u'RowList.exe', u'TrayMenu.exe'])
             self.assertEqual(files_list.item('RebarTest.exe').window_text(), 'RebarTest.exe')
 
         def test_set_backend_to_window_uia(self):

pywinauto/unittests/test_uiawrapper.py

@@ -204,9 +204,9 @@ def test_move_window(self):
             )
             time.sleep(0.1)
             logger = ActionLogger()
-            logger.log("prev_rect = ", prev_rect)
-            logger.log("new_rect = ", new_rect)
-            logger.log("self.dlg.rectangle() = ", self.dlg.rectangle())
+            logger.log("prev_rect = %s", prev_rect)
+            logger.log("new_rect = %s", new_rect)
+            logger.log("self.dlg.rectangle() = %s", self.dlg.rectangle())
             self.assertEqual(self.dlg.rectangle(), new_rect)
 
             self.dlg.move_window(prev_rect)

pywinauto/unittests/test_win32controls.py

@@ -462,9 +462,9 @@ def tearDown(self):
         finally:
             self.app.kill()
 
-    def test_print_control_identifiers(self):
-        """Test that print_control_identifiers() doesn't crash with the non-English characters"""
-        self.dlg.print_control_identifiers()
+    def test_dump_tree(self):
+        """Test that dump_tree() doesn't crash with the non-English characters"""
+        self.dlg.dump_tree()
 
     def test_set_text(self):
         """Test setting the text of the edit control"""

pywinauto/windows/uia_element_info.py

@@ -416,12 +416,15 @@ def enabled(self):
     @property
     def rectangle(self):
         """Return rectangle of the element"""
-        bound_rect = self._element.CurrentBoundingRectangle
         rect = RECT()
-        rect.left = bound_rect.left
-        rect.top = bound_rect.top
-        rect.right = bound_rect.right
-        rect.bottom = bound_rect.bottom
+        try:
+            bound_rect = self._element.CurrentBoundingRectangle
+            rect.left = bound_rect.left
+            rect.top = bound_rect.top
+            rect.right = bound_rect.right
+            rect.bottom = bound_rect.bottom
+        except COMError:
+            pass
         return rect
 
     def dump_window(self):