Merge pull request #15 from gmatteo/master

Add Command class and stream_has_colours

docs/monty.rst

@@ -115,6 +115,14 @@ monty.string module
     :undoc-members:
     :show-inheritance:
 
+monty.subprocess module
+-------------------
+
+.. automodule:: monty.subprocess
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 monty.tempfile module
 ---------------------
 

monty/collections.py

@@ -7,6 +7,8 @@
 __email__ = 'ongsp@ucsd.edu'
 __date__ = '1/24/14'
 
+import collections
+
 
 class frozendict(dict):
     """
@@ -45,12 +47,10 @@ class AttrDict(dict):
     to the traditional way obj['foo']"
 
     Example:
-        >> d = AttrDict(foo=1, bar=2)
-        >> d["foo"] == d.foo
-        True
-        >> d.bar = "hello"
-        >> d.bar
-        'hello'
+        >>> d = AttrDict(foo=1, bar=2)
+        >>> assert d["foo"] == d.foo
+        >>> d.bar = "hello"
+        >>> assert d.bar == "hello"
     """
     def __init__(self, *args, **kwargs):
         super(AttrDict, self).__init__(*args, **kwargs)
@@ -84,3 +84,58 @@ def __getattribute__(self, name):
     def __setattr__(self, name, value):
         raise KeyError("You cannot modify attribute %s of %s" % (name, self.__class__.__name__))
 
+
+class MongoDict(object):
+    """
+    This dict-like object allows one to access the entries in a nested dict as attributes. 
+    Entries (attributes) cannot be modified. It also provides Ipython tab completion hence this object 
+    is particularly useful if you need to analyze a nested dict interactively (e.g. documents
+    extracted from a MongoDB database).
+
+    >>> m = MongoDict({'a': {'b': 1}, 'x': 2}) 
+    >>> assert m.a.b == 1 and m.x == 2
+    >>> assert "a" in m and "b" in m.a
+
+    .. note:: 
+
+        Cannot inherit from ABC collections.Mapping because otherwise 
+        dict.keys and dict.items will pollute the namespace.
+        e.g MongoDict({"keys": 1}).keys would be the ABC dict method.
+    """
+    def __init__(self, *args, **kwargs):
+        self.__dict__["_mongo_dict_"] = dict(*args, **kwargs)
+
+    def __repr__(self):
+        return str(self)
+
+    def __str__(self):
+        return str(self._mongo_dict_)
+
+    def __setattr__(self, name, value):
+        raise NotImplementedError("You cannot modify attribute %s of %s" % (name, self.__class__.__name__))
+
+    def __getattribute__(self, name):
+        try:
+            return super(MongoDict, self).__getattribute__(name)
+        except:
+            #raise
+            try:
+                a = self._mongo_dict_[name]
+                if isinstance(a, collections.Mapping):
+                    a = self.__class__(a)
+                return a
+            except Exception as exc:
+                raise AttributeError(str(exc))
+
+    def __getitem__(self, slice):
+        return self._mongo_dict_.__getitem__(slice)
+
+    def __iter__(self):
+        return iter(self._mongo_dict_)
+
+    def __len__(self):
+        return len(self._mongo_dict_)
+
+    def __dir__(self):
+        """For Ipython tab completion. See http://ipython.org/ipython-doc/dev/config/integrating.html"""
+        return sorted(list(k for k in self._mongo_dict_ if not callable(k)))

monty/shutil.py

@@ -123,4 +123,5 @@ def decompress_dir(path):
     """
     for parent, subdirs, files in os.walk(path):
         for f in files:
-            decompress_file(os.path.join(parent, f))
+            decompress_file(os.path.join(parent, f))
+

monty/subprocess.py

@@ -0,0 +1,92 @@
+# coding: utf-8
+from __future__ import absolute_import, print_function, division, unicode_literals
+
+__author__ = 'Matteo Giantomass'
+__copyright__ = "Copyright 2014, The Materials Virtual Lab"
+__version__ = '0.1'
+__maintainer__ = 'Matteo Giantomassi'
+__email__ = 'gmatteo@gmail.com'
+__date__ = '10/26/14'
+
+
+class Command(object):
+    """
+    Enables to run subprocess commands in a different thread with TIMEOUT option.
+
+    Based on jcollado's solution:
+        http://stackoverflow.com/questions/1191374/subprocess-with-timeout/4825933#4825933
+    and
+        https://gist.github.com/kirpit/1306188
+
+    .. attribute:: retcode
+        
+        Return code of the subprocess
+
+    .. attribute:: killed
+
+        True if subprocess has been killed due to the timeout
+
+    .. attribute:: output
+
+        stdout of the subprocess
+
+    .. attribute:: error
+
+        stderr of the subprocess
+
+    Example:
+        com = Command("sleep 1").run(timeout=2)
+        print(com.retcode, com.killed, com.output, com.output)
+    """
+    def __init__(self, command):
+        from .string import is_string
+        if is_string(command):
+            import shlex
+            command = shlex.split(command)
+
+        self.command = command
+        self.process = None
+        self.retcode = None
+        self.output, self.error = '', ''
+        self.killed = False
+
+    def __str__(self):
+        return "command: %s, retcode: %s" % (str(self.command), str(self.retcode))
+
+    def run(self, timeout=None, **kwargs):
+        """
+        Run a command in a separated thread and wait timeout seconds.
+        kwargs are keyword arguments passed to Popen.
+
+        Return: self
+        """
+        from subprocess import Popen, PIPE
+        def target(**kwargs):
+            try:
+                #print('Thread started')
+                self.process = Popen(self.command, **kwargs)
+                self.output, self.error = self.process.communicate()
+                self.retcode = self.process.returncode
+                #print('Thread stopped')
+            except:
+                import traceback
+                self.error = traceback.format_exc()
+                self.retcode = -1
+
+        # default stdout and stderr
+        if 'stdout' not in kwargs: kwargs['stdout'] = PIPE
+        if 'stderr' not in kwargs: kwargs['stderr'] = PIPE
+
+        # thread
+        import threading
+        thread = threading.Thread(target=target, kwargs=kwargs)
+        thread.start()
+        thread.join(timeout)
+
+        if thread.is_alive():
+            #print("Terminating process")
+            self.process.terminate()
+            self.killed = True
+            thread.join()
+
+        return self

monty/termcolor.py

@@ -181,3 +181,19 @@ def cprint(text, color=None, on_color=None, attrs=None, **kwargs):
             ['underline'])
     cprint('Reversed green on red color', 'green', 'on_red', ['reverse'])
 
+
+def stream_has_colours(stream):
+    """
+    True if stream supports colours. Python cookbook, #475186
+    """
+    if not hasattr(stream, "isatty"):
+        return False
+
+    if not stream.isatty():
+        return False  # auto color only on TTYs
+    try:
+        import curses
+        curses.setupterm()
+        return curses.tigetnum("colors") > 2
+    except:
+        return False  # guess false in case of error

tests/test_shutil.py

@@ -81,5 +81,6 @@ def test_compress_and_decompress_file(self):
     def tearDown(self):
         os.remove(os.path.join(test_dir, "tempfile"))
 
+
 if __name__ == "__main__":
     unittest.main()

tests/test_subprocess.py

@@ -0,0 +1,21 @@
+import unittest
+
+from monty.subprocess import Command
+
+class CommandTest(unittest.TestCase):
+    def test_command(self):
+        """Test Command class"""
+        sleep05 = Command("sleep 0.5")
+
+        sleep05.run(timeout=1)
+        print(sleep05)
+        self.assertEqual(sleep05.retcode, 0)
+        self.assertFalse(sleep05.killed)
+
+        sleep05.run(timeout=0.1)
+        self.assertNotEqual(sleep05.retcode, 0)
+        self.assertTrue(sleep05.killed)
+
+
+if __name__ == "__main__":
+    unittest.main()