fix select: to more CSS style selection

[htmltree-py/htmltree.git] / htmltree.py

diff --git a/htmltree.py b/htmltree.py
index 4e50bb3..597da01 100644 (file)
--- a/htmltree.py
+++ b/htmltree.py
@@ -168,19 +168,21 @@ class HTMLElement(list):
         if len(terms) == 0:
             return []
         results = self
-        # at first, select #id
-        ids = [x[1:] for x in terms if x[0] == "#"]
-        if len(ids) != 0:
-            results = [results.get_element_by_id(ids[-1]),]
-        # next, select .class
-        cs = [x[1:] for x in terms if x[0] == "."]
-        t = []
-        if len(cs) != 0:
-            for c in cs:
-                t.extend(results.get_elements_by_class(c))
+        for pat in terms:
+            t = []
+            for elem in results:
+                t.extend(self._select_pattern(pat, elem))
             results = t
         return results
 
+    def _select_pattern(self, pat, elem):
+        results = []
+        if pat[0] == "#":
+            results = [elem.get_element_by_id(pat[1:]),]
+        elif pat[0] == ".":
+            results = elem.get_elements_by_class(pat[1:])
+        return [x for x in results if x]
+
     def select_1st(self, expr):
         r = self.select(expr)
         if len(r) == 0: