diff --git a/docs/monty.rst b/docs/monty.rst index 96b9a3870..9cbcd61f4 100644 --- a/docs/monty.rst +++ b/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 --------------------- diff --git a/monty/collections.py b/monty/collections.py index b705eb630..69a6bf805 100644 --- a/monty/collections.py +++ b/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))) diff --git a/monty/shutil.py b/monty/shutil.py index 71a948633..ce8b61e7e 100644 --- a/monty/shutil.py +++ b/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)) \ No newline at end of file + decompress_file(os.path.join(parent, f)) + diff --git a/monty/subprocess.py b/monty/subprocess.py new file mode 100644 index 000000000..5ec72dae1 --- /dev/null +++ b/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 diff --git a/monty/termcolor.py b/monty/termcolor.py index 84f32753a..c13a60b97 100644 --- a/monty/termcolor.py +++ b/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 diff --git a/tests/test_shutil.py b/tests/test_shutil.py index 782489741..2e3f84ee2 100644 --- a/tests/test_shutil.py +++ b/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() diff --git a/tests/test_subprocess.py b/tests/test_subprocess.py new file mode 100644 index 000000000..53e793ed7 --- /dev/null +++ b/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()