From 19d2d56e36db4fd0bc32d662e1c9f9c08aba1494 Mon Sep 17 00:00:00 2001
From: Takeshi KOMIYA <i.tkomiya@gmail.com>
Date: Tue, 21 Feb 2017 00:17:38 +0900
Subject: [PATCH] Fix IndexBuilder.prune() and load() supports _filenames field

---
 sphinx/search/__init__.py |  4 ++++
 tests/test_search.py      | 28 ++++++++++++++++++++++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/sphinx/search/__init__.py b/sphinx/search/__init__.py
index beffe2549..4b309be2e 100644
--- a/sphinx/search/__init__.py
+++ b/sphinx/search/__init__.py
@@ -271,6 +271,7 @@ class IndexBuilder(object):
            frozen.get('envversion') != self.env.version:
             raise ValueError('old format')
         index2fn = frozen['docnames']
+        self._filenames = dict(zip(index2fn, frozen['filenames']))
         self._titles = dict(zip(index2fn, frozen['titles']))
 
         def load_terms(mapping):
@@ -361,10 +362,13 @@ class IndexBuilder(object):
     def prune(self, docnames):
         """Remove data for all docnames not in the list."""
         new_titles = {}
+        new_filenames = {}
         for docname in docnames:
             if docname in self._titles:
                 new_titles[docname] = self._titles[docname]
+                new_filenames[docname] = self._filenames[docname]
         self._titles = new_titles
+        self._filenames = new_filenames
         for wordnames in itervalues(self._mapping):
             wordnames.intersection_update(docnames)
         for wordnames in itervalues(self._title_mapping):
diff --git a/tests/test_search.py b/tests/test_search.py
index b9f0b9338..19b322dd8 100644
--- a/tests/test_search.py
+++ b/tests/test_search.py
@@ -184,16 +184,21 @@ def test_IndexBuilder():
     index2.load(stream, 'pickle')
 
     assert index2._titles == index._titles
-    assert index2._filenames == {}
+    assert index2._filenames == index._filenames
     assert index2._mapping == index._mapping
     assert index2._title_mapping == index._title_mapping
     assert index2._objtypes == {}
     assert index2._objnames == {}
 
+    # freeze after load
+    assert index2.freeze() == index.freeze()
+    assert index2._objtypes == index._objtypes
+    assert index2._objnames == index._objnames
+
     # prune
     index.prune(['docname2'])
     assert index._titles == {'docname2': 'title2'}
-    assert index._filenames == {'docname': 'filename', 'docname2': 'filename2'}
+    assert index._filenames == {'docname2': 'filename2'}
     assert index._mapping == {
         'fermion': {'docname2'},
         'comment': {'docname2'},
@@ -204,3 +209,22 @@ def test_IndexBuilder():
     assert index._title_mapping == {'section_titl': {'docname2'}}
     assert index._objtypes == {('dummy', 'objtype'): 0}
     assert index._objnames == {0: ('dummy', 'objtype', 'objtype')}
+
+    # freeze after prune
+    assert index.freeze() == {
+        'docnames': ('docname2',),
+        'envversion': '1.0',
+        'filenames': ['filename2'],
+        'objects': {},
+        'objnames': {0: ('dummy', 'objtype', 'objtype')},
+        'objtypes': {0: 'dummy:objtype'},
+        'terms': {'comment': 0,
+                  'fermion': 0,
+                  'index': 0,
+                  'non': 0,
+                  'test': 0},
+        'titles': ('title2',),
+        'titleterms': {'section_titl': 0}
+    }
+    assert index._objtypes == {('dummy', 'objtype'): 0}
+    assert index._objnames == {0: ('dummy', 'objtype', 'objtype')}