From d940d44f08ceb5ada49f4ba6d9128523aecdf552 Mon Sep 17 00:00:00 2001
From: Steven Myint <git@stevenmyint.com>
Date: Thu, 4 Jan 2018 15:32:10 -0800
Subject: [PATCH] Narrow how much we filter out complex cases

---
 autoflake.py      | 45 +++++++++++++++++++++++++++++----------------
 test_autoflake.py | 24 ++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 16 deletions(-)

diff --git a/autoflake.py b/autoflake.py
index dafae07..46d1d93 100755
--- a/autoflake.py
+++ b/autoflake.py
@@ -151,25 +151,38 @@ def duplicate_key_line_numbers(messages, source):
         if isinstance(message, pyflakes.messages.MultiValueRepeatedKeyLiteral)]
 
     if messages:
-        good = True
         # Filter out complex cases. We don't want to bother trying to parse
-        # this stuff and get it right.
+        # this stuff and get it right. We can do it on a key-by-key basis.
+
+        key_to_messages = create_key_to_messages_dict(messages)
+
         lines = source.split('\n')
-        for message in messages:
-            line = lines[message.lineno - 1]
-            key = message.message_args[0]
-
-            if (
-                line.rstrip().endswith((':', '\\')) or
-                not dict_entry_has_key(line, key) or
-                '#' in line or
-                not line.rstrip().endswith(',')
-            ):
-                good = False
-
-        if good:
+
+        for (key, messages) in key_to_messages.items():
+            good = True
             for message in messages:
-                yield message.lineno
+                line = lines[message.lineno - 1]
+                key = message.message_args[0]
+
+                if (
+                    line.rstrip().endswith((':', '\\')) or
+                    not dict_entry_has_key(line, key) or
+                    '#' in line or
+                    not line.rstrip().endswith(',')
+                ):
+                    good = False
+
+            if good:
+                for message in messages:
+                    yield message.lineno
+
+
+def create_key_to_messages_dict(messages):
+    """Return dict mapping the key to list of messages."""
+    dictionary = collections.defaultdict(lambda: [])
+    for message in messages:
+        dictionary[message.message_args[0]].append(message)
+    return dictionary
 
 
 def check(source):
diff --git a/test_autoflake.py b/test_autoflake.py
index 66d6611..2e922dc 100755
--- a/test_autoflake.py
+++ b/test_autoflake.py
@@ -464,6 +464,30 @@ def test_filter_code_should_ignore_complex_case_of_duplicate_key(self):
             ''.join(autoflake.filter_code(code,
                                           remove_duplicate_keys=True)))
 
+    def test_filter_code_should_ignore_complex_case_of_duplicate_key_partially(
+            self):
+        """We only handle simple cases."""
+        code = """\
+a = {(0,1): 1, (0, 1): 'two',
+  (0,1): 3,
+  (2,3): 4,
+  (2,3): 4,
+  (2,3): 5,
+}
+print(a)
+"""
+
+        self.assertEqual(
+            """\
+a = {(0,1): 1, (0, 1): 'two',
+  (0,1): 3,
+  (2,3): 5,
+}
+print(a)
+""",
+            ''.join(autoflake.filter_code(code,
+                                          remove_duplicate_keys=True)))
+
     def test_filter_code_should_ignore_more_cases_of_duplicate_key(self):
         """We only handle simple cases."""
         code = """\