Fix and move NamespaceTrie to spack.util.naming

- fix up routines in namespace trie.
- trie can now hold intermediate elements.
- trie now has a test case.

4 files changed, 154 insertions, 42 deletions

diff --git a/lib/spack/spack/packages.py b/lib/spack/spack/packages.py
index 6005523bc0..3a74ad6790 100644
--- a/lib/spack/spack/packages.py
+++ b/lib/spack/spack/packages.py
@@ -58,45 +58,6 @@ def _autospec(function):
     return converter
 
 
-class NamespaceTrie(object):
-    def __init__(self):
-        self._elements = {}
-
-
-    def __setitem__(self, namespace, repo):
-        parts = namespace.split('.')
-        cur = self._elements
-        for p in parts[:-1]:
-            if p not in cur:
-                cur[p] = {}
-            cur = cur[p]
-
-        cur[parts[-1]] = repo
-
-
-    def __getitem__(self, namespace):
-        parts = namespace.split('.')
-        cur = self._elements
-        for p in parts:
-            if p not in cur:
-                raise KeyError("Can't find namespace %s in trie" % namespace)
-            cur = cur[p]
-        return cur
-
-
-    def __contains__(self, namespace):
-        parts = namespace.split('.')
-        cur = self._elements
-        for p in parts:
-            if not isinstance(cur, dict):
-                return False
-            if p not in cur:
-                return False
-            cur  = cur[p]
-        return True
-
-
-
 class PackageFinder(object):
     """A PackageFinder is a wrapper around a list of PackageDBs.
 
@@ -172,7 +133,7 @@ class PackageFinder(object):
 
 
     def providers_for(self, vpkg_name):
-        # TODO: THIS IS WRONG; shoudl use more than first repo
+        # TODO: THIS IS WRONG; should use more than first repo
         return self.repos[0].providers_for(vpkg_name)
 
 
diff --git a/lib/spack/spack/test/__init__.py b/lib/spack/spack/test/__init__.py
index 0f776bfea4..620d1fd362 100644
--- a/lib/spack/spack/test/__init__.py
+++ b/lib/spack/spack/test/__init__.py
@@ -59,7 +59,8 @@ test_names = ['versions',
               'configure_guess',
               'unit_install',
               'lock',
-              'database']
+              'database',
+              'namespace_trie']
 
 
 def list_tests():
diff --git a/lib/spack/spack/test/namespace_trie.py b/lib/spack/spack/test/namespace_trie.py
new file mode 100644
index 0000000000..191abbe9e6
--- /dev/null
+++ b/lib/spack/spack/test/namespace_trie.py
@@ -0,0 +1,83 @@
+##############################################################################
+# Copyright (c) 2013, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+#
+# This file is part of Spack.
+# Written by Todd Gamblin, tgamblin@llnl.gov, All rights reserved.
+# LLNL-CODE-647188
+#
+# For details, see https://scalability-llnl.github.io/spack
+# Please also see the LICENSE file for our notice and the LGPL.
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License (as published by
+# the Free Software Foundation) version 2.1 dated February 1999.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the IMPLIED WARRANTY OF
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the terms and
+# conditions of the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program; if not, write to the Free Software Foundation,
+# Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
+##############################################################################
+import unittest
+from spack.util.naming import NamespaceTrie
+
+
+class NamespaceTrieTest(unittest.TestCase):
+
+    def setUp(self):
+        self.trie = NamespaceTrie()
+
+
+    def test_add_single(self):
+        self.trie['foo'] = 'bar'
+        self.assertEqual(self.trie['foo'], 'bar')
+        self.assertTrue('foo' in self.trie)
+
+
+    def test_add_multiple(self):
+        self.trie['foo.bar'] = 'baz'
+        self.assertEqual(self.trie['foo.bar'], 'baz')
+
+        self.assertFalse('foo' in self.trie)
+        self.assertFalse('foo.bar.baz' in self.trie)
+        self.assertTrue('foo.bar' in self.trie)
+
+
+    def test_add_three(self):
+        # add a three-level namespace
+        self.trie['foo.bar.baz'] = 'quux'
+        self.assertEqual(self.trie['foo.bar.baz'], 'quux')
+
+        self.assertFalse('foo' in self.trie)
+        self.assertFalse('foo.bar' in self.trie)
+        self.assertTrue('foo.bar.baz' in self.trie)
+
+        # Try to add a second element in a higher space
+        self.trie['foo.bar'] = 'blah'
+
+        self.assertFalse('foo' in self.trie)
+
+        self.assertTrue('foo.bar' in self.trie)
+        self.assertEqual(self.trie['foo.bar'], 'blah')
+
+        self.assertTrue('foo.bar.baz' in self.trie)
+        self.assertEqual(self.trie['foo.bar.baz'], 'quux')
+
+
+    def test_add_none_single(self):
+        self.trie['foo'] = None
+        self.assertEqual(self.trie['foo'], None)
+        self.assertTrue('foo' in self.trie)
+
+
+    def test_add_none_multiple(self):
+        self.trie['foo.bar'] = None
+        self.assertEqual(self.trie['foo.bar'], None)
+
+        self.assertFalse('foo' in self.trie)
+        self.assertFalse('foo.bar.baz' in self.trie)
+        self.assertTrue('foo.bar' in self.trie)
diff --git a/lib/spack/spack/util/naming.py b/lib/spack/spack/util/naming.py
index a7b6e2b436..475062bb38 100644
--- a/lib/spack/spack/util/naming.py
+++ b/lib/spack/spack/util/naming.py
@@ -3,11 +3,12 @@ from __future__ import absolute_import
 import string
 import itertools
 import re
+from StringIO import StringIO
 
 import spack
 
 __all__ = ['mod_to_class', 'spack_module_to_python_module', 'valid_module_name',
-           'validate_module_name', 'possible_spack_module_names']
+           'validate_module_name', 'possible_spack_module_names', 'NamespaceTrie']
 
 # Valid module names can contain '-' but can't start with it.
 _valid_module_re = r'^\w[\w-]*$'
@@ -90,3 +91,69 @@ class InvalidModuleNameError(spack.error.SpackError):
         super(InvalidModuleNameError, self).__init__(
             "Invalid module name: " + name)
         self.name = name
+
+
+class NamespaceTrie(object):
+    class Element(object):
+        def __init__(self, value):
+            self.value = value
+
+
+    def __init__(self, separator='.'):
+        self._subspaces = {}
+        self._value = None
+        self._sep = separator
+
+
+    def __setitem__(self, namespace, value):
+        first, sep, rest = namespace.partition(self._sep)
+
+        if not first:
+            self._value = NamespaceTrie.Element(value)
+            return
+
+        if first not in self._subspaces:
+            self._subspaces[first] = NamespaceTrie()
+
+        self._subspaces[first][rest] = value
+
+
+    def _get_helper(self, namespace, full_name):
+        first, sep, rest = namespace.partition(self._sep)
+        if not first:
+            if not self._value:
+                raise KeyError("Can't find namespace '%s' in trie" % full_name)
+            return self._value.value
+        elif first not in self._subspaces:
+            raise KeyError("Can't find namespace '%s' in trie" % full_name)
+        else:
+            return self._subspaces[first]._get_helper(rest, full_name)
+
+
+    def __getitem__(self, namespace):
+        return self._get_helper(namespace, namespace)
+
+
+    def __contains__(self, namespace):
+        first, sep, rest = namespace.partition(self._sep)
+        if not first:
+            return self._value is not None
+        elif first not in self._subspaces:
+            return False
+        else:
+            return rest in self._subspaces[first]
+
+
+    def _str_helper(self, stream, level=0):
+        indent = (level * '    ')
+        for name in sorted(self._subspaces):
+            stream.write(indent + name + '\n')
+            if self._value:
+                stream.write(indent + '  ' + repr(self._value.value))
+            stream.write(self._subspaces[name]._str_helper(stream, level+1))
+
+
+    def __str__(self):
+        stream = StringIO()
+        self._str_helper(stream)
+        return stream.getvalue()