Add optional usage of CreateTreeWalker to iter_children and iter_desc…

…endants methods

pywinauto/unittests/test_uia_element_info.py

@@ -12,6 +12,7 @@
 
 if UIA_support:
     from pywinauto.windows.uia_element_info import UIAElementInfo
+    from pywinauto.windows.uia_defines import IUIA
 
 mfc_samples_folder = os.path.join(
     os.path.dirname(__file__), r"..\..\apps\WPF_samples")
@@ -106,32 +107,58 @@ def test_descendants_generator(self):
             descendants = [desc for desc in self.ctrl.iter_descendants(depth=3)]
             self.assertSequenceEqual(self.ctrl.descendants(depth=3), descendants)
 
-        def test_use_findall_children(self):
-            """Test use FindAll option for children method"""
+        def test_use_property_conditions_children(self):
+            """Test use CondTreeWalker option for children method"""
             with mock.patch.object(self.ctrl._element, 'FindAll', wraps=self.ctrl._element.FindAll) as mock_findall:
-                self.assertEqual(UIAElementInfo.use_findall, False)
+                self.assertEqual(UIAElementInfo.use_property_conditions, False)
 
-                UIAElementInfo.use_findall = True
+                UIAElementInfo.use_property_conditions = True
                 self.assertEqual(len(self.ctrl.children()), 5)
                 self.assertEqual(mock_findall.call_count, 1)
 
-                UIAElementInfo.use_findall = False
+                UIAElementInfo.use_property_conditions = False
                 self.assertEqual(len(self.ctrl.children()), 5)
                 self.assertEqual(mock_findall.call_count, 1)
 
-        def test_use_findall_descendants(self):
-            """Test use FindAll option for descendants method"""
+        def test_use_property_conditions_descendants(self):
+            """Test use CondTreeWalker option for descendants method"""
             with mock.patch.object(self.ctrl._element, 'FindAll', wraps=self.ctrl._element.FindAll) as mock_findall:
-                self.assertEqual(UIAElementInfo.use_findall, False)
+                self.assertEqual(UIAElementInfo.use_property_conditions, False)
 
-                UIAElementInfo.use_findall = True
+                UIAElementInfo.use_property_conditions = True
                 self.assertEqual(len(self.ctrl.descendants(depth=1)), 5)
                 self.assertEqual(mock_findall.call_count, 1)
 
-                UIAElementInfo.use_findall = False
+                UIAElementInfo.use_property_conditions = False
                 self.assertEqual(len(self.ctrl.descendants(depth=1)), 5)
                 self.assertEqual(mock_findall.call_count, 1)
 
+        def test_use_property_conditions_iter_children(self):
+            """Test use CondTreeWalker option for iter_children method"""
+            with mock.patch.object(IUIA().iuia, 'CreateTreeWalker', wraps=IUIA().iuia.CreateTreeWalker) as mock_walker:
+                self.assertEqual(UIAElementInfo.use_property_conditions, False)
+
+                UIAElementInfo.use_property_conditions = True
+                self.assertSequenceEqual(self.ctrl.children(), list(self.ctrl.iter_children()))
+                self.assertEqual(mock_walker.call_count, 1)
+
+                UIAElementInfo.use_property_conditions = False
+                self.assertSequenceEqual(self.ctrl.children(), list(self.ctrl.iter_children()))
+                self.assertEqual(mock_walker.call_count, 1)
+
+        def test_use_property_conditions_iter_descendants(self):
+            """Test use CondTreeWalker option for descendants method"""
+            with mock.patch.object(IUIA().iuia, 'CreateTreeWalker', wraps=IUIA().iuia.CreateTreeWalker) as mock_walker:
+                self.assertEqual(UIAElementInfo.use_property_conditions, False)
+
+                UIAElementInfo.use_property_conditions = True
+                self.assertSequenceEqual(self.ctrl.descendants(), list(self.ctrl.iter_descendants()))
+                self.assertEqual(mock_walker.call_count, 1)
+
+                UIAElementInfo.use_property_conditions = False
+                self.assertSequenceEqual(self.ctrl.descendants(), list(self.ctrl.iter_descendants()))
+                self.assertEqual(mock_walker.call_count, 1)
+
 if __name__ == "__main__":
     if UIA_support:
         unittest.main()

pywinauto/windows/uia_element_info.py

@@ -96,7 +96,7 @@ class UIAElementInfo(ElementInfo):
         "top_level_only": ("depth", {True: 1, False: None}),
     }
 
-    use_findall = False
+    use_property_conditions = False
 
     def __init__(self, handle_or_elem=None, cache_enable=False):
         """
@@ -331,9 +331,17 @@ def iter_children(self, **kwargs):
            class_name, control_type, content_only and/or title.
         """
         cache_enable = kwargs.pop('cache_enable', False)
-        for element in self._iter_children_raw():
-            if is_element_satisfying_criteria(element, **kwargs):
-                yield UIAElementInfo(element, cache_enable)
+        if UIAElementInfo.use_property_conditions:
+            cond = IUIA().build_condition(**kwargs)
+            tree_walker = IUIA().iuia.CreateTreeWalker(cond)
+            element = tree_walker.GetFirstChildElement(self._element)
+            while element:
+                yield UIAElementInfo(element)
+                element = tree_walker.GetNextSiblingElement(element)
+        else:
+            for element in self._iter_children_raw():
+                if is_element_satisfying_criteria(element, **kwargs):
+                    yield UIAElementInfo(element, cache_enable)
 
     def iter_descendants(self, **kwargs):
         """Iterate over descendants of the element"""
@@ -344,22 +352,30 @@ def iter_descendants(self, **kwargs):
 
         if depth == 0:
             return
-        for child in self._iter_children_raw():
-            if is_element_satisfying_criteria(child, **kwargs):
-                yield UIAElementInfo(child, cache_enable)
-            if depth is not None:
-                kwargs["depth"] = depth - 1
-            for c in UIAElementInfo(child, cache_enable).iter_descendants(**kwargs):
-                if is_element_satisfying_criteria(c._element, **kwargs):
+        if UIAElementInfo.use_property_conditions:
+            for child in self.iter_children(**kwargs):
+                yield child
+                if depth is not None:
+                    kwargs["depth"] = depth - 1
+                for c in child.iter_descendants(**kwargs):
                     yield c
+        else:
+            for child in self._iter_children_raw():
+                if is_element_satisfying_criteria(child, **kwargs):
+                    yield UIAElementInfo(child, cache_enable)
+                if depth is not None:
+                    kwargs["depth"] = depth - 1
+                for c in UIAElementInfo(child, cache_enable).iter_descendants(**kwargs):
+                    if is_element_satisfying_criteria(c._element, **kwargs):
+                        yield c
 
     def children(self, **kwargs):
         """Return a list of only immediate children of the element
 
          * **kwargs** is a criteria to reduce a list by process,
            class_name, control_type, content_only and/or title.
         """
-        if UIAElementInfo.use_findall:
+        if UIAElementInfo.use_property_conditions:
             cache_enable = kwargs.pop('cache_enable', False)
             cond = IUIA().build_condition(**kwargs)
             return self._get_elements(IUIA().tree_scope["children"], cond, cache_enable)
@@ -372,7 +388,7 @@ def descendants(self, **kwargs):
          * **kwargs** is a criteria to reduce a list by process,
            class_name, control_type, content_only and/or title.
         """
-        if UIAElementInfo.use_findall:
+        if UIAElementInfo.use_property_conditions:
             cache_enable = kwargs.pop('cache_enable', False)
             depth = kwargs.pop('depth', None)
             cond = IUIA().build_condition(**kwargs)