regex: fix attribute handling in make_pattern

Needs more testing.

1 files changed, 29 insertions, 3 deletions

diff --git a/regex/__init__.py b/regex/__init__.py
index 2a123a6..3738251 100644
--- a/regex/__init__.py
+++ b/regex/__init__.py
@@ -54,7 +54,11 @@ def make_pattern(tree, selected):
             return Tree(label)
         elif label == 'Call':
             # omit arguments, keep the original function name
-            return Tree(label, [tree[0]])
+            if tree[0][0].label == 'Attribute':
+                # method call
+                return Tree(label, [Tree('func', [Tree('Attribute', tree[0][0][1:])])])
+            else:
+                return Tree(label, [tree[0]])
         else:
             return tree
 
@@ -67,9 +71,31 @@ def make_pattern(tree, selected):
     # some selected nodes
 
     # extra nodes to include in patterns, depending on node label
-    if label in {'Call', 'FunctionDef'}:
-        # store function name
+    if label == 'FunctionDef':
+        # include function name
         pat = Tree(label, [tree[0]] + [s for s in subpats if s is not None])
+    elif label == 'Attribute':
+        # include attribute name
+        children = []
+        if subpats[0] is not None:
+            children += [subpats[0]]
+        children += tree[1:]
+        pat = Tree(label, children)
+    elif label == 'Call':
+        # include function/method name
+        children = []
+        if subpats[0] is not None:
+            children += [subpats[0]]
+        else:
+            func = tree[0]
+            if func[0].label == 'Attribute':
+                # method call: include attribute (method) name
+                children += [Tree('func', [Tree('Attribute', func[0][1:])])]
+            else:
+                # function call: include function name
+                children += [func]
+        children += [s for s in subpats[1:] if s is not None]
+        pat = Tree(label, children)
     elif label == 'args':
         # (args arg1 … argN)
         # include at least an 'arg' node in the pattern for each