<commit_before>from io import StringIO
import pytest

from rich.console import Console
from rich.text import Span, Text
from rich.measure import Measurement

def test_span():
    span = Span(1, 10, "foo")
    repr(span)
    assert bool(span)
    assert not Span(10, 10, "foo")

def test_span_split():
    assert Span(5, 10, "foo").split(2) == (Span(5, 10, "foo"), None)
    assert Span(5, 10, "foo").split(15) == (Span(5, 10, "foo"), None)
    assert Span(0, 10, "foo").split(5) == (Span(0, 5, "foo"), Span(5, 10, "foo"))

def test_span_move():
    assert Span(5, 10, "foo").move(2) == Span(7, 12, "foo")

def test_span_right_crop():
    assert Span(5, 10, "foo").right_crop(15) == Span(5, 10, "foo")
    assert Span(5, 10, "foo").right_crop(7) == Span(5, 7, "foo")

def test_len():
    assert len(Text("foo")) == 3

def test_bool():
    assert Text("foo")
    assert not Text("")

def test_str():
    assert str(Text("foo")) == "foo"

def test_repr():
    assert isinstance(repr(Text("foo")), str)

def test_add():
    text = Text("foo") + Text("bar")
    assert str(text) == "foobar"
    assert Text("foo").__add__(1) == NotImplemented

def test_eq():
    assert Text("foo") == Text("foo")
    assert Text("foo") != Text("bar")
    assert Text("foo").__eq__(1) == NotImplemented

def test_text_property():
    text = Text("foo")
    text.append("bar")
    text.append("baz")
    assert text.text == "foobarbaz"

def test_text_property_setter():
    test = Text("foo")
    test.text = "bar"
    assert str(test) == "bar"
    test = Text()
    test.append("Hello, World", "bold")
    test.text = "Hello"
    assert str(test) == "Hello"
    assert test._spans == [Span(0, 5, "bold")]

def test_from_markup():
    text = Text.from_markup("Hello, [bold]World![/bold]")
    assert str(text) == "Hello, World!"
    assert text._spans == [Span(7, 13, "bold")]

def test_copy():
    test = Text()
    test.append("Hello", "bold")
    test.append(" ")
    test.append("World", "italic")
    test_copy = test.copy()
    assert test == test_copy
    assert test is not test_copy

def test_rstrip():
    test = Text("Hello, World!    ")
    test.rstrip()
    assert str(test) == "Hello, World!"

def test_stylize():
    test = Text("Hello, World!")
    test.stylize(7, 11, "bold")
    assert test._spans == [Span(7, 11, "bold")]
    test.stylize(20, 25, "bold")
    assert test._spans == [Span(7, 11, "bold")]

def test_set_length():
    test = Text("Hello")
    test.set_length(5)
    assert test == Text("Hello")

    test = Text("Hello")
    test.set_length(10)
    assert test == Text("Hello     ")

    test = Text("Hello World")
    test.stylize(0, 5, "bold")
    test.stylize(7, 9, "italic")

    test.set_length(3)
    expected = Text()
    expected.append("Hel", "bold")
    assert test == expected

def test_console_width():
    console = Console()
    test = Text("Hello World!\nfoobarbaz")
    assert test.__measure__(console, 80) == Measurement(9, 12)
    assert Text(" " * 4).__measure__(console, 80) == Measurement(4, 4)

def test_join():
    test = Text("bar").join([Text("foo", "red"), Text("baz", "blue")])
    assert str(test) == "foobarbaz"
    assert test._spans == [Span(0, 3, "red"), Span(3, 6, ""), Span(6, 9, "blue")]

def test_trim_spans():
    test = Text("Hello")
    test._spans[:] = [Span(0, 3, "red"), Span(3, 6, "green"), Span(6, 9, "blue")]
    test._trim_spans()
    assert test._spans == [Span(0, 3, "red"), Span(3, 5, "green")]

def test_pad_left():
    test = Text("foo")
    test.pad_left(3, "X")
    assert str(test) == "XXXfoo"

def test_pad_right():
    test = Text("foo")
    test.pad_right(3, "X")
    assert str(test) == "fooXXX"

def test_append():
    test = Text("foo")
    test.append("bar")
    assert str(test) == "foobar"
    test.append(Text("baz", "bold"))
    assert str(test) == "foobarbaz"
    assert test._spans == [Span(6, 9, "bold")]

    with pytest.raises(ValueError):
        test.append(Text("foo"), "bar")

    with pytest.raises(TypeError):
        test.append(1)

def test_split():
    test = Text()
    test.append("foo", "red")
    test.append("\n")
    test.append("bar", "green")
    test.append("\n")

    line1 = Text()
    line1.append("foo", "red")
    line2 = Text()
    line2.append("bar", "green")
    split = test.split("\n")
    assert len(split) == 2
    assert split[0] == line1
    assert split[1] == line2

    assert list(Text("foo").split("\n")) == [Text("foo")]

def test_divide():
    lines = Text("foo").divide([])
    assert len(lines) == 1
    assert lines[0] == Text("foo")

    text = Text()
    text.append("foo", "bold")
    lines = text.divide([1, 2])
    assert len(lines) == 3
    assert str(lines[0]) == "f"
    assert str(lines[1]) == "o"
    assert str(lines[2]) == "o"
    assert lines[0]._spans == [Span(0, 1, "bold")]
    assert lines[1]._spans == [Span(0, 1, "bold")]
    assert lines[2]._spans == [Span(0, 1, "bold")]

    text = Text()
    text.append("foo", "red")
    text.append("bar", "green")
    text.append("baz", "blue")
    lines = text.divide([8])
    assert len(lines) == 2
    assert str(lines[0]) == "foobarba"
    assert str(lines[1]) == "z"
    assert lines[0]._spans == [
        Span(0, 3, "red"),
        Span(3, 6, "green"),
        Span(6, 8, "blue"),
    ]
    assert lines[1]._spans == [Span(0, 1, "blue")]

    lines = text.divide([1])
    assert len(lines) == 2
    assert str(lines[0]) == "f"
    assert str(lines[1]) == "oobarbaz"
    assert lines[0]._spans == [Span(0, 1, "red")]
    assert lines[1]._spans == [
        Span(0, 2, "red"),
        Span(2, 5, "green"),
        Span(5, 8, "blue"),
    ]

def test_right_crop():
    test = Text()
    test.append("foobar", "red")
    test.right_crop(3)
    assert str(test) == "foo"
    assert test._spans == [Span(0, 3, "red")]

def test_wrap_4():
    test = Text("foo bar baz")
    lines = test.wrap(4)
    assert len(lines) == 3
    assert lines[0] == Text("foo ")
    assert lines[1] == Text("bar ")
    assert lines[2] == Text("baz ")

def test_wrap_3():
    test = Text("foo bar baz")
    lines = test.wrap(3)
    assert len(lines) == 3
    assert lines[0] == Text("foo")
    assert lines[1] == Text("bar")
    assert lines[2] == Text("baz")

def test_wrap_long():
    test = Text("abracadabra")
    lines = test.wrap(4)
    assert len(lines) == 3
    assert lines[0] == Text("abra")
    assert lines[1] == Text("cada")
    assert lines[2] == Text("bra ")

def test_wrap_long_words():
    test = Text("X 123456789")
    lines = test.wrap(4)
    print(lines)
    assert len(lines) == 3
    assert lines[0] == Text("X 12")
    assert lines[1] == Text("3456")
    assert lines[2] == Text("789 ")

def test_fit():
    test = Text("Hello\nWorld")
    lines = test.fit(3)
    assert str(lines[0]) == "Hel"
    assert str(lines[1]) == "Wor"

def test_render():
    console = Console(width=15, record=True)
    test = Text.from_markup(
        "[u][b]Where[/b] there is a [i]Will[/i], there is a Way.[/u]"
    )
    console.print(test)
    output = console.export_text(styles=True)
    print(repr(output))
    expected = "\x1b[1;4mWhere\x1b[0m\x1b[4m there is \x1b[0m\n\x1b[4ma \x1b[0m\x1b[3;4mWill\x1b[0m\x1b[4m, there \x1b[0m\n\x1b[4mis a Way.\x1b[0m\n"
    assert output == expected

@pytest.mark.parametrize(
    "print_text,result",
    [
        (("."), ".\n"),
        ((".", "."), ". .\n"),
        (("Hello", "World", "!"), "Hello World !\n"),
    ],
)
def test_print(print_text, result):
    console = Console(record=True)
    console.print(*print_text)
    assert console.export_text(styles=False) == result

@pytest.mark.parametrize(
    "print_text,result",
    [(("."), ".X"), ((".", "."), "..X"), (("Hello", "World", "!"), "HelloWorld!X"),],
)
def test_print_sep_end(print_text, result):
    console = Console(record=True, file=StringIO())
    console.print(*print_text, sep="", end="X")
    assert console.file.getvalue() == result

def test_tabs_to_spaces():
    test = Text("\tHello\tWorld", tab_size=8)
    assert test.tabs_to_spaces().text == "        Hello   World"

    test = Text("\tHello\tWorld", tab_size=4)
    assert test.tabs_to_spaces().text == "    Hello   World"

    test = Text(".\t..\t...\t....\t", tab_size=4)
    assert test.tabs_to_spaces().text == ".   ..  ... ....    "

def test_markup_switch():
    """Test markup can be disabled."""
    console = Console(file=StringIO(), markup=False)
    console.print("[bold]foo[/bold]")
    assert console.file.getvalue() == "[bold]foo[/bold]\n"

def test_emoji():
    """Test printing emoji codes."""
    console = Console(file=StringIO())
    console.print(":+1:")
    assert console.file.getvalue() == "👍\n"

def test_emoji_switch():
    """Test emoji can be disabled."""
    console = Console(file=StringIO(), emoji=False)
    console.print(":+1:")
    assert console.file.getvalue() == ":+1:\n"
<commit_msg>tests for Text<commit_after>from io import StringIO
import pytest

from rich.console import Console
from rich.text import Span, Text
from rich.measure import Measurement

def test_span():
    span = Span(1, 10, "foo")
    repr(span)
    assert bool(span)
    assert not Span(10, 10, "foo")

def test_span_split():
    assert Span(5, 10, "foo").split(2) == (Span(5, 10, "foo"), None)
    assert Span(5, 10, "foo").split(15) == (Span(5, 10, "foo"), None)
    assert Span(0, 10, "foo").split(5) == (Span(0, 5, "foo"), Span(5, 10, "foo"))

def test_span_move():
    assert Span(5, 10, "foo").move(2) == Span(7, 12, "foo")

def test_span_right_crop():
    assert Span(5, 10, "foo").right_crop(15) == Span(5, 10, "foo")
    assert Span(5, 10, "foo").right_crop(7) == Span(5, 7, "foo")

def test_len():
    assert len(Text("foo")) == 3

def test_bool():
    assert Text("foo")
    assert not Text("")

def test_str():
    assert str(Text("foo")) == "foo"

def test_repr():
    assert isinstance(repr(Text("foo")), str)

def test_add():
    text = Text("foo") + Text("bar")
    assert str(text) == "foobar"
    assert Text("foo").__add__(1) == NotImplemented

def test_eq():
    assert Text("foo") == Text("foo")
    assert Text("foo") != Text("bar")
    assert Text("foo").__eq__(1) == NotImplemented

def test_contain():
    test = Text("foobar")
    assert "foo" in test
    assert "foo " not in test
    assert Text("bar") in test
    assert None not in test

def test_text_property():
    text = Text("foo")
    text.append("bar")
    text.append("baz")
    assert text.text == "foobarbaz"

def test_text_property_setter():
    test = Text("foo")
    test.text = "bar"
    assert str(test) == "bar"
    test = Text()
    test.append("Hello, World", "bold")
    test.text = "Hello"
    assert str(test) == "Hello"
    assert test._spans == [Span(0, 5, "bold")]

def test_from_markup():
    text = Text.from_markup("Hello, [bold]World![/bold]")
    assert str(text) == "Hello, World!"
    assert text._spans == [Span(7, 13, "bold")]

def test_copy():
    test = Text()
    test.append("Hello", "bold")
    test.append(" ")
    test.append("World", "italic")
    test_copy = test.copy()
    assert test == test_copy
    assert test is not test_copy

def test_rstrip():
    test = Text("Hello, World!    ")
    test.rstrip()
    assert str(test) == "Hello, World!"

def test_stylize():
    test = Text("Hello, World!")
    test.stylize(7, 11, "bold")
    assert test._spans == [Span(7, 11, "bold")]
    test.stylize(20, 25, "bold")
    assert test._spans == [Span(7, 11, "bold")]

def test_highlight_regex():
    test = Text("peek-a-boo")

    count = test.highlight_regex(r"NEVER_MATCH", "red")
    assert count == 0
    assert len(test._spans) == 0

    # text: peek-a-boo
    # indx: 0123456789
    count = test.highlight_regex(r"[a|e|o]+", "red")
    assert count == 3
    assert sorted(test._spans) == [
        Span(1, 3, "red"),
        Span(5, 6, "red"),
        Span(8, 10, "red"),
    ]

    test = Text("Ada Lovelace, Alan Turing")
    count = test.highlight_regex(
        r"(?P<yellow>[A-Za-z]+)[ ]+(?P<red>[A-Za-z]+)(?P<NEVER_MATCH>NEVER_MATCH)*"
    )

    # The number of matched name should be 2
    assert count == 2
    assert sorted(test._spans) == [
        Span(0, 3, "yellow"),  # Ada
        Span(4, 12, "red"),  # Lovelace
        Span(14, 18, "yellow"),  # Alan
        Span(19, 25, "red"),  # Turing
    ]

def test_highlight_words():
    test = Text("Do NOT! touch anything!")
    words = ["NOT", "!"]
    count = test.highlight_words(words, "red")
    assert count == 3
    assert sorted(test._spans) == [
        Span(3, 6, "red"),  # NOT
        Span(6, 7, "red"),  # !
        Span(22, 23, "red"),  # !
    ]

    # regex escape test
    test = Text("[o|u]aeiou")
    words = ["[a|e|i]", "[o|u]"]
    count = test.highlight_words(words, "red")
    assert count == 1
    assert test._spans == [Span(0, 5, "red")]

    # case sensitive
    test = Text("AB Ab aB ab")
    words = ["AB"]

    count = test.highlight_words(words, "red")
    assert count == 1
    assert test._spans == [Span(0, 2, "red")]

    test = Text("AB Ab aB ab")
    count = test.highlight_words(words, "red", False)
    assert count == 4

def test_set_length():
    test = Text("Hello")
    test.set_length(5)
    assert test == Text("Hello")

    test = Text("Hello")
    test.set_length(10)
    assert test == Text("Hello     ")

    test = Text("Hello World")
    test.stylize(0, 5, "bold")
    test.stylize(7, 9, "italic")

    test.set_length(3)
    expected = Text()
    expected.append("Hel", "bold")
    assert test == expected

def test_console_width():
    console = Console()
    test = Text("Hello World!\nfoobarbaz")
    assert test.__measure__(console, 80) == Measurement(9, 12)
    assert Text(" " * 4).__measure__(console, 80) == Measurement(4, 4)

def test_join():
    test = Text("bar").join([Text("foo", "red"), Text("baz", "blue")])
    assert str(test) == "foobarbaz"
    assert test._spans == [Span(0, 3, "red"), Span(3, 6, ""), Span(6, 9, "blue")]

def test_trim_spans():
    test = Text("Hello")
    test._spans[:] = [Span(0, 3, "red"), Span(3, 6, "green"), Span(6, 9, "blue")]
    test._trim_spans()
    assert test._spans == [Span(0, 3, "red"), Span(3, 5, "green")]

def test_pad_left():
    test = Text("foo")
    test.pad_left(3, "X")
    assert str(test) == "XXXfoo"

def test_pad_right():
    test = Text("foo")
    test.pad_right(3, "X")
    assert str(test) == "fooXXX"

def test_append():
    test = Text("foo")
    test.append("bar")
    assert str(test) == "foobar"
    test.append(Text("baz", "bold"))
    assert str(test) == "foobarbaz"
    assert test._spans == [Span(6, 9, "bold")]

    with pytest.raises(ValueError):
        test.append(Text("foo"), "bar")

    with pytest.raises(TypeError):
        test.append(1)

def test_split():
    test = Text()
    test.append("foo", "red")
    test.append("\n")
    test.append("bar", "green")
    test.append("\n")

    line1 = Text()
    line1.append("foo", "red")
    line2 = Text()
    line2.append("bar", "green")
    split = test.split("\n")
    assert len(split) == 2
    assert split[0] == line1
    assert split[1] == line2

    assert list(Text("foo").split("\n")) == [Text("foo")]

def test_divide():
    lines = Text("foo").divide([])
    assert len(lines) == 1
    assert lines[0] == Text("foo")

    text = Text()
    text.append("foo", "bold")
    lines = text.divide([1, 2])
    assert len(lines) == 3
    assert str(lines[0]) == "f"
    assert str(lines[1]) == "o"
    assert str(lines[2]) == "o"
    assert lines[0]._spans == [Span(0, 1, "bold")]
    assert lines[1]._spans == [Span(0, 1, "bold")]
    assert lines[2]._spans == [Span(0, 1, "bold")]

    text = Text()
    text.append("foo", "red")
    text.append("bar", "green")
    text.append("baz", "blue")
    lines = text.divide([8])
    assert len(lines) == 2
    assert str(lines[0]) == "foobarba"
    assert str(lines[1]) == "z"
    assert lines[0]._spans == [
        Span(0, 3, "red"),
        Span(3, 6, "green"),
        Span(6, 8, "blue"),
    ]
    assert lines[1]._spans == [Span(0, 1, "blue")]

    lines = text.divide([1])
    assert len(lines) == 2
    assert str(lines[0]) == "f"
    assert str(lines[1]) == "oobarbaz"
    assert lines[0]._spans == [Span(0, 1, "red")]
    assert lines[1]._spans == [
        Span(0, 2, "red"),
        Span(2, 5, "green"),
        Span(5, 8, "blue"),
    ]

def test_right_crop():
    test = Text()
    test.append("foobar", "red")
    test.right_crop(3)
    assert str(test) == "foo"
    assert test._spans == [Span(0, 3, "red")]

def test_wrap_4():
    test = Text("foo bar baz")
    lines = test.wrap(4)
    assert len(lines) == 3
    assert lines[0] == Text("foo ")
    assert lines[1] == Text("bar ")
    assert lines[2] == Text("baz ")

def test_wrap_3():
    test = Text("foo bar baz")
    lines = test.wrap(3)
    assert len(lines) == 3
    assert lines[0] == Text("foo")
    assert lines[1] == Text("bar")
    assert lines[2] == Text("baz")

def test_wrap_long():
    test = Text("abracadabra")
    lines = test.wrap(4)
    assert len(lines) == 3
    assert lines[0] == Text("abra")
    assert lines[1] == Text("cada")
    assert lines[2] == Text("bra ")

def test_wrap_long_words():
    test = Text("X 123456789")
    lines = test.wrap(4)
    print(lines)
    assert len(lines) == 3
    assert lines[0] == Text("X 12")
    assert lines[1] == Text("3456")
    assert lines[2] == Text("789 ")

def test_fit():
    test = Text("Hello\nWorld")
    lines = test.fit(3)
    assert str(lines[0]) == "Hel"
    assert str(lines[1]) == "Wor"

def test_render():
    console = Console(width=15, record=True)
    test = Text.from_markup(
        "[u][b]Where[/b] there is a [i]Will[/i], there is a Way.[/u]"
    )
    console.print(test)
    output = console.export_text(styles=True)
    print(repr(output))
    expected = "\x1b[1;4mWhere\x1b[0m\x1b[4m there is \x1b[0m\n\x1b[4ma \x1b[0m\x1b[3;4mWill\x1b[0m\x1b[4m, there \x1b[0m\n\x1b[4mis a Way.\x1b[0m\n"
    assert output == expected

@pytest.mark.parametrize(
    "print_text,result",
    [
        (("."), ".\n"),
        ((".", "."), ". .\n"),
        (("Hello", "World", "!"), "Hello World !\n"),
    ],
)
def test_print(print_text, result):
    console = Console(record=True)
    console.print(*print_text)
    assert console.export_text(styles=False) == result

@pytest.mark.parametrize(
    "print_text,result",
    [(("."), ".X"), ((".", "."), "..X"), (("Hello", "World", "!"), "HelloWorld!X"),],
)
def test_print_sep_end(print_text, result):
    console = Console(record=True, file=StringIO())
    console.print(*print_text, sep="", end="X")
    assert console.file.getvalue() == result

def test_tabs_to_spaces():
    test = Text("\tHello\tWorld", tab_size=8)
    assert test.tabs_to_spaces().text == "        Hello   World"

    test = Text("\tHello\tWorld", tab_size=4)
    assert test.tabs_to_spaces().text == "    Hello   World"

    test = Text(".\t..\t...\t....\t", tab_size=4)
    assert test.tabs_to_spaces().text == ".   ..  ... ....    "

    test = Text("No Tabs")
    assert test.tabs_to_spaces().text == "No Tabs"

def test_markup_switch():
    """Test markup can be disabled."""
    console = Console(file=StringIO(), markup=False)
    console.print("[bold]foo[/bold]")
    assert console.file.getvalue() == "[bold]foo[/bold]\n"

def test_emoji():
    """Test printing emoji codes."""
    console = Console(file=StringIO())
    console.print(":+1:")
    assert console.file.getvalue() == "👍\n"

def test_emoji_switch():
    """Test emoji can be disabled."""
    console = Console(file=StringIO(), emoji=False)
    console.print(":+1:")
    assert console.file.getvalue() == ":+1:\n"
<|endoftext|>


<commit_before>""" A parser and other parser related classes. """

import pyparsing

from .models import Command, Input, Output, Grammar
from .models import command_template_factory as ctf

class Parser(object):

    def __init__(self, string):
        self.string = string
        self.commands = []
        self.paths = []
        self.files = []

    def consume(self):
        """ Converts the lexer tokens into valid statements. This process
        also checks command syntax.
        """
        grammar = Grammar.overall
        grammar.ignore(Grammar.comment)
        first_pass = grammar.parseString(self.string)
        lowered = { key.lower(): val for key, val in first_pass.iteritems() }

        try:
            txt_files = [''.join(f) for f in lowered['files'].asList()]
        except KeyError:
            txt_files = []
        try:
            txt_paths = [''.join(p) for p in lowered['paths'].asList()]
        except KeyError:
            txt_paths = []
        try:
            txt_commands = [''.join(c) for c in lowered['commands'].asList()]
        except KeyError:
            txt_commands = []

        for c in txt_commands:
            try:
                self.commands.append(Grammar.command.parseString(c))
            except pyparsing.ParseException as e:
                raise ValueError('Invalid command. Verify line {} is '
                    'correct.\n{}'.format(e.lineno, c))

        for f in txt_files:
            try:
                self.files.append(Grammar.file.parseString(f))
            except pyparsing.ParseException as e:
                raise ValueError('Invalid file. Verify line {} is '
                    'correct.\n{}'.format(e.lineno, f))

        for p in txt_paths:
            try:
                self.paths.append(Grammar.path.parseString(p))
            except pyparsing.ParseException as e:
                raise ValueError('Invalid path. Verify line {} is '
                    'correct.\n{}'.format(e.lineno, p))

        self.paths = ctf.get_paths(self.paths)
        self.files = ctf.get_files(self.files)

        self.paths.reverse()
        self.files.reverse()
        self.commands.reverse()

        return ctf.get_command_templates(self.commands, self.files[:],
            self.paths[:])
<commit_msg>Simplifies parser logic<commit_after>""" A parser and other parser related classes. """

import pyparsing

from .models import Command, Input, Output, Grammar
from .models import command_template_factory as ctf

class Parser(object):

    def __init__(self, string):
        self.string = string
        self.commands = []
        self.paths = []
        self.files = []

    def consume(self):
        """ Converts the lexer tokens into valid statements. This process
        also checks command syntax.
        """
        grammar = Grammar.overall
        grammar.ignore(Grammar.comment)
        first_pass = grammar.parseString(self.string)
        lowered = { key.lower(): val for key, val in first_pass.iteritems() }

        self.files = self._get('files', lowered, Grammar.file)
        self.paths = self._get('paths', lowered, Grammar.path)
        self.commands = self._get('commands', lowered, Grammar.command)

        self.paths = ctf.get_paths(self.paths)
        self.files = ctf.get_files(self.files)

        self.paths.reverse()
        self.files.reverse()
        self.commands.reverse()

        return ctf.get_command_templates(self.commands, self.files[:],
            self.paths[:])

    def _get(self, type, parser_result, grammar):
        """ Given a type and a dict of parser results, parse the overall results
        using the more detailed parse grammar.
        """
        try:
            txt_lines = [''.join(f) for f in parser_result[type].asList()]
        except KeyError:
            txt_lines = []

        results = []
        for c in txt_lines:
            try:
                results.append(grammar.parseString(c))
            except pyparsing.ParseException as e:
                raise ValueError('Invalid {}. Verify line {} is '
                    'correct.\n{}'.format(type, e.lineno, c))
        return results
<|endoftext|>


<commit_before>#!/usr/bin/env python

'''
This module defines the VaspInputSet abstract base class and
a concrete implementation for the Materials Project.  The basic
concept behind an input set is to specify a scheme to generate
a consistent set of Vasp inputs from a structure without further
user intervention. This ensures comparability across runs.
'''

from __future__ import division

__author__="Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Nov 16, 2011"

import os
import abc
import ConfigParser
import json

from pymatgen.io.vaspio import Incar, Poscar, Potcar, Kpoints

class AbstractVaspInputSet(object):
    """
    Abstract base class representing a set of Vasp input parameters.
    The idea is that using a VaspInputSet, a complete set of input files (INPUT, KPOINTS, POSCAR and POTCAR)
    can be generated in an automated fashion for any structure.
    """
    __metaclass__ = abc.ABCMeta

    def get_poscar(self, structure):
        '''
        Returns Poscar from a structure.
        '''
        return Poscar(structure)

    @abc.abstractmethod
    def get_kpoints(self, structure):
        '''
        Returns Kpoints from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    @abc.abstractmethod
    def get_incar(self, structure):
        '''
        Returns Incar from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    @abc.abstractmethod
    def get_potcar(self, structure):
        '''
        Returns Potcar from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    @abc.abstractmethod
    def get_potcar_symbols(self, structure):
        '''
        Returns Potcar from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    def get_all_vasp_input(self, structure, generate_potcar = True):
        '''
        Returns all input files as a dict of {filename: vaspio object}

        Arguments:
            structure:
                Structure object
            generate_potcar:
                Set to False to generate a POTCAR.spec file instead of a POTCAR,
                which contains the POTCAR labels but not the actual POTCAR. Defaults
                to True.

        Returns:
            dict of {filename: file_as_string}, e.g., {'INCAR':'EDIFF=1e-4...'}
        '''
        if generate_potcar:
            return {'INCAR':self.get_incar(structure), 'KPOINTS':self.get_kpoints(structure),
                    'POSCAR': self.get_poscar(structure), 'POTCAR': self.get_potcar(structure)}
        else:
            return {'INCAR':self.get_incar(structure), 'KPOINTS':self.get_kpoints(structure),
                    'POSCAR': self.get_poscar(structure), 'POTCAR.spec': "\n".join(self.get_potcar_symbols(structure))}

    def write_input(self, structure, output_dir, make_dir_if_not_present = True):
        """
        Writes a set of VASP input to a directory.

        Arguments:
            structure:
                Structure object
            output_dir:
                Directory to output the VASP input files
            make_dir_if_not_present:
                Set to True if you want the directory (and the whole path) to be created if it is not present.
        """
        if make_dir_if_not_present and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        for k,v in self.get_all_vasp_input(structure).items():
            v.write_file(os.path.join(output_dir, k))

class VaspInputSet(AbstractVaspInputSet):
    """
    Standard implementation of VaspInputSet, which can be extended by specific
    implementations.
    """

    def __init__(self, name):
        self.name = name
        module_dir = os.path.dirname(os.path.abspath(__file__))
        self._config = ConfigParser.SafeConfigParser()
        self._config.optionxform = str
        self._config.readfp(open(os.path.join(module_dir, "VaspInputSets.cfg")))
        self.potcar_settings = dict(self._config.items(self.name + 'POTCAR'))
        self.kpoints_settings = dict(self._config.items(self.name + 'KPOINTS'))
        self.incar_settings = dict(self._config.items(self.name+'INCAR'))
        for key in ['MAGMOM', 'LDAUU', 'LDAUJ', 'LDAUL']:
            if key in self.incar_settings:
                self.incar_settings[key] = json.loads(self.incar_settings[key])

    def get_incar(self, structure):
        incar = Incar()
        poscar = Poscar(structure)
        for key, setting in self.incar_settings.items():
            if key == "MAGMOM":
                incar[key] =  [setting.get(site.specie.symbol, 0.6) for site in structure]
            elif key in ['LDAUU', 'LDAUJ', 'LDAUL']:
                incar[key] =  [setting.get(sym, 0) for sym in poscar.site_symbols]
            elif key == "EDIFF":
                incar[key] =  float(setting) * structure.num_sites
            else:
                incar[key] = setting

        has_u = ("LDAUU" in incar and sum(incar['LDAUU']) > 0)
        if not has_u:
            for key in incar.keys():
                if key.startswith('LDAU'):
                    del incar[key]

        return incar

    #get_poscar method inherited from AbstractVaspInputSet

    def get_potcar(self, structure):
        return Potcar(self.get_potcar_symbols(structure))

    def get_potcar_symbols(self, structure):
        p = self.get_poscar(structure)
        elements = p.site_symbols
        potcar_symbols = []
        for el in elements:
            potcar_symbols.append(self.potcar_settings[el] if el in self.potcar_settings else el)
        return potcar_symbols

    def get_kpoints(self, structure):
        '''
        Writes out a KPOINTS file using the fully automated grid method. Uses Gamma centered meshes
        for hexagonal cells and Monkhorst-Pack grids otherwise.

        Algorithm:
            Uses a simple approach scaling the number of divisions along each
            reciprocal lattice vector proportional to its length.
        '''

        latt = structure.lattice
        lengths = latt.abc
        ngrid = int(self.kpoints_settings['grid_density']) / structure.num_sites

        mult = (ngrid*lengths[0]*lengths[1] * lengths[2]) ** (1/3)

        num_div = [int(round(1/lengths[i] * mult)) for i in xrange(3)]
        #ensure that numDiv[i] > 0

        num_div = [i if i > 0 else 1 for i in num_div]

        angles = latt.angles
        hex_angle_tol = 5 #in degrees
        hex_length_tol = 0.01 #in angstroms
        right_angles = [i for i in xrange(3) if abs(angles[i] - 90) < hex_angle_tol]
        hex_angles = [i for i in xrange(3) if abs(angles[i] - 60) < hex_angle_tol or abs(angles[i] - 120) < hex_angle_tol]

        is_hexagonal = (len(right_angles) == 2 and len(hex_angles) == 1 and abs(lengths[right_angles[0]] == lengths[right_angles[1]]) < hex_length_tol)

        style = 'Gamma'
        if not is_hexagonal:
            num_div = [i + i % 2 for i in num_div]
            style = 'Monk'
        comment = "pymatgen generated Materials Project kpoints with grid density = " + self.kpoints_settings['grid_density'] + ' per atom.'
        num_kpts = 0
        return Kpoints(comment, num_kpts, style, [num_div], [0,0,0])

    def __str__(self):
        output = [self.name]
        output.append("")
        section_names = ['INCAR settings', 'KPOINTS settings', 'POTCAR settings']
        count = 0
        for d in [self.incar_settings, self.kpoints_settings, self.potcar_settings]:
            output.append(section_names[count])
            for k, v in d.items():
                output.append("%s = %s" % (k, str(v)))
            output.append("")
            count += 1

        return "\n".join(output)

class MITVaspInputSet(VaspInputSet):
    """
    Standard implementation of VaspInputSet utilizing parameters in the MIT
    High-throughput project.
    The parameters are chosen specifically for a high-throughput project,
    which means in general smaller pseudopotentials were chosen.

    Please refer to A Jain, G. Hautier, C. Moore, S. P. Ong, C. Fischer,
    T. Mueller, K. A. Persson, G. Ceder (2011).
    A high-throughput infrastructure for density functional theory calculations.
    Computational Materials Science, 50(8), 2295-2310.
    doi:10.1016/j.commatsci.2011.02.023 for more information.
    """
    def __init__(self):
        super(MITVaspInputSet, self).__init__("MITMatgen")

class MITHSEVaspInputSet(VaspInputSet):
    """
    Typical implementation of input set for a HSE run.
    """
    def __init__(self):
        super(MITHSEVaspInputSet, self).__init__("MITHSE")

class MaterialsProjectVaspInputSet(VaspInputSet):
    """
    Implementation of VaspInputSet utilizing parameters in the public
    Materials Project. Typically, the psuedopotentials chosen contain more
    electrons than the MIT parameters, and the k-point grid is ~50% more dense.
    The LDAUU parameters are also different due to the different psps used,
    which result in different fitted values (even though the methodology of
    fitting is exactly the same as the MIT scheme).
    """
    def __init__(self):
        super(MaterialsProjectVaspInputSet, self).__init__("MaterialsProject")
<commit_msg>Fixed bug in kpoint style assignment using Enum<commit_after>#!/usr/bin/env python

'''
This module defines the VaspInputSet abstract base class and
a concrete implementation for the Materials Project.  The basic
concept behind an input set is to specify a scheme to generate
a consistent set of Vasp inputs from a structure without further
user intervention. This ensures comparability across runs.
'''

from __future__ import division

__author__="Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Nov 16, 2011"

import os
import abc
import ConfigParser
import json

from pymatgen.io.vaspio import Incar, Poscar, Potcar, Kpoints

class AbstractVaspInputSet(object):
    """
    Abstract base class representing a set of Vasp input parameters.
    The idea is that using a VaspInputSet, a complete set of input files (INPUT, KPOINTS, POSCAR and POTCAR)
    can be generated in an automated fashion for any structure.
    """
    __metaclass__ = abc.ABCMeta

    def get_poscar(self, structure):
        '''
        Returns Poscar from a structure.
        '''
        return Poscar(structure)

    @abc.abstractmethod
    def get_kpoints(self, structure):
        '''
        Returns Kpoints from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    @abc.abstractmethod
    def get_incar(self, structure):
        '''
        Returns Incar from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    @abc.abstractmethod
    def get_potcar(self, structure):
        '''
        Returns Potcar from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    @abc.abstractmethod
    def get_potcar_symbols(self, structure):
        '''
        Returns Potcar from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    def get_all_vasp_input(self, structure, generate_potcar = True):
        '''
        Returns all input files as a dict of {filename: vaspio object}

        Arguments:
            structure:
                Structure object
            generate_potcar:
                Set to False to generate a POTCAR.spec file instead of a POTCAR,
                which contains the POTCAR labels but not the actual POTCAR. Defaults
                to True.

        Returns:
            dict of {filename: file_as_string}, e.g., {'INCAR':'EDIFF=1e-4...'}
        '''
        if generate_potcar:
            return {'INCAR':self.get_incar(structure), 'KPOINTS':self.get_kpoints(structure),
                    'POSCAR': self.get_poscar(structure), 'POTCAR': self.get_potcar(structure)}
        else:
            return {'INCAR':self.get_incar(structure), 'KPOINTS':self.get_kpoints(structure),
                    'POSCAR': self.get_poscar(structure), 'POTCAR.spec': "\n".join(self.get_potcar_symbols(structure))}

    def write_input(self, structure, output_dir, make_dir_if_not_present = True):
        """
        Writes a set of VASP input to a directory.

        Arguments:
            structure:
                Structure object
            output_dir:
                Directory to output the VASP input files
            make_dir_if_not_present:
                Set to True if you want the directory (and the whole path) to be created if it is not present.
        """
        if make_dir_if_not_present and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        for k,v in self.get_all_vasp_input(structure).items():
            v.write_file(os.path.join(output_dir, k))

class VaspInputSet(AbstractVaspInputSet):
    """
    Standard implementation of VaspInputSet, which can be extended by specific
    implementations.
    """

    def __init__(self, name):
        self.name = name
        module_dir = os.path.dirname(os.path.abspath(__file__))
        self._config = ConfigParser.SafeConfigParser()
        self._config.optionxform = str
        self._config.readfp(open(os.path.join(module_dir, "VaspInputSets.cfg")))
        self.potcar_settings = dict(self._config.items(self.name + 'POTCAR'))
        self.kpoints_settings = dict(self._config.items(self.name + 'KPOINTS'))
        self.incar_settings = dict(self._config.items(self.name+'INCAR'))
        for key in ['MAGMOM', 'LDAUU', 'LDAUJ', 'LDAUL']:
            if key in self.incar_settings:
                self.incar_settings[key] = json.loads(self.incar_settings[key])

    def get_incar(self, structure):
        incar = Incar()
        poscar = Poscar(structure)
        for key, setting in self.incar_settings.items():
            if key == "MAGMOM":
                incar[key] =  [setting.get(site.specie.symbol, 0.6) for site in structure]
            elif key in ['LDAUU', 'LDAUJ', 'LDAUL']:
                incar[key] =  [setting.get(sym, 0) for sym in poscar.site_symbols]
            elif key == "EDIFF":
                incar[key] =  float(setting) * structure.num_sites
            else:
                incar[key] = setting

        has_u = ("LDAUU" in incar and sum(incar['LDAUU']) > 0)
        if not has_u:
            for key in incar.keys():
                if key.startswith('LDAU'):
                    del incar[key]

        return incar

    #get_poscar method inherited from AbstractVaspInputSet

    def get_potcar(self, structure):
        return Potcar(self.get_potcar_symbols(structure))

    def get_potcar_symbols(self, structure):
        p = self.get_poscar(structure)
        elements = p.site_symbols
        potcar_symbols = []
        for el in elements:
            potcar_symbols.append(self.potcar_settings[el] if el in self.potcar_settings else el)
        return potcar_symbols

    def get_kpoints(self, structure):
        '''
        Writes out a KPOINTS file using the fully automated grid method. Uses Gamma centered meshes
        for hexagonal cells and Monkhorst-Pack grids otherwise.

        Algorithm:
            Uses a simple approach scaling the number of divisions along each
            reciprocal lattice vector proportional to its length.
        '''

        latt = structure.lattice
        lengths = latt.abc
        ngrid = int(self.kpoints_settings['grid_density']) / structure.num_sites

        mult = (ngrid*lengths[0]*lengths[1] * lengths[2]) ** (1/3)

        num_div = [int(round(1/lengths[i] * mult)) for i in xrange(3)]
        #ensure that numDiv[i] > 0

        num_div = [i if i > 0 else 1 for i in num_div]

        angles = latt.angles
        hex_angle_tol = 5 #in degrees
        hex_length_tol = 0.01 #in angstroms
        right_angles = [i for i in xrange(3) if abs(angles[i] - 90) < hex_angle_tol]
        hex_angles = [i for i in xrange(3) if abs(angles[i] - 60) < hex_angle_tol or abs(angles[i] - 120) < hex_angle_tol]

        is_hexagonal = (len(right_angles) == 2 and len(hex_angles) == 1 and abs(lengths[right_angles[0]] == lengths[right_angles[1]]) < hex_length_tol)

        style = Kpoints.supported_modes.Gamma
        if not is_hexagonal:
            num_div = [i + i % 2 for i in num_div]
            style = Kpoints.supported_modes.Monkhorst
        comment = "pymatgen generated Materials Project kpoints with grid density = " + self.kpoints_settings['grid_density'] + ' per atom.'
        num_kpts = 0
        return Kpoints(comment, num_kpts, style, [num_div], [0,0,0])

    def __str__(self):
        output = [self.name]
        output.append("")
        section_names = ['INCAR settings', 'KPOINTS settings', 'POTCAR settings']
        count = 0
        for d in [self.incar_settings, self.kpoints_settings, self.potcar_settings]:
            output.append(section_names[count])
            for k, v in d.items():
                output.append("%s = %s" % (k, str(v)))
            output.append("")
            count += 1

        return "\n".join(output)

class MITVaspInputSet(VaspInputSet):
    """
    Standard implementation of VaspInputSet utilizing parameters in the MIT
    High-throughput project.
    The parameters are chosen specifically for a high-throughput project,
    which means in general smaller pseudopotentials were chosen.

    Please refer to A Jain, G. Hautier, C. Moore, S. P. Ong, C. Fischer,
    T. Mueller, K. A. Persson, G. Ceder (2011).
    A high-throughput infrastructure for density functional theory calculations.
    Computational Materials Science, 50(8), 2295-2310.
    doi:10.1016/j.commatsci.2011.02.023 for more information.
    """
    def __init__(self):
        super(MITVaspInputSet, self).__init__("MITMatgen")

class MITHSEVaspInputSet(VaspInputSet):
    """
    Typical implementation of input set for a HSE run.
    """
    def __init__(self):
        super(MITHSEVaspInputSet, self).__init__("MITHSE")

class MaterialsProjectVaspInputSet(VaspInputSet):
    """
    Implementation of VaspInputSet utilizing parameters in the public
    Materials Project. Typically, the psuedopotentials chosen contain more
    electrons than the MIT parameters, and the k-point grid is ~50% more dense.
    The LDAUU parameters are also different due to the different psps used,
    which result in different fitted values (even though the methodology of
    fitting is exactly the same as the MIT scheme).
    """
    def __init__(self):
        super(MaterialsProjectVaspInputSet, self).__init__("MaterialsProject")
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# The MIT License (MIT)
#
# Copyright (c) 2014 William T. James for Storj Labs
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import unittest
import os
import redis
import hashlib
import subprocess
import binascii

import RandomIO
from sys import platform as _platform

if _platform.startswith('linux') or _platform == 'darwin':
    cat_cmd = 'cat'
    iotools_call = ['IOTools.py']
elif _platform == 'win32':
    cat_cmd = 'type'
    iotools_call = ['python', os.path.join('../bin', 'IOTools.py')]

class TestRandomIO(unittest.TestCase):

    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_gen(self):
        s = RandomIO.RandomIO()
        b = s.read(100)

        self.assertEqual(len(b), 100)

        self.assertEqual(
            RandomIO.RandomIO(123456).read(100),
            RandomIO.RandomIO(123456).read(100))
        self.assertEqual(RandomIO.RandomIO(b'byte string seed').read(
            100), RandomIO.RandomIO(b'byte string seed').read(100))
        self.assertEqual(RandomIO.RandomIO(1.23456).read(
            100), RandomIO.RandomIO(1.23456).read(100))

    def test_consistent(self):
        s1 = RandomIO.RandomIO('seed string')
        s2 = RandomIO.RandomIO('seed string')
        s3 = RandomIO.RandomIO('seed string')
        s4 = RandomIO.RandomIO('another seed')

        self.assertEqual(s1.read(100), s2.read(100))
        self.assertNotEqual(s3.read(100), s4.read(100))

    def test_crossplatform(self):
        string_seed1 = b'\t\xb0\xef\xd9\x05p\xe1W\x17\x8a9\xc6!;^6\x1d\xadj\
\xb4#n\x1d/\x12+\xe6\xb1\x80\xc86\x06I\xc4!\x8b39\x84E\x1d\x14\xdf\x14e\x12\
\xfa\xf0\r\x1b'

        s = RandomIO.RandomIO('seed1').read(50)

        self.assertEqual(s, string_seed1)

        string_123456 = b'\x18\xb2\xce\x8a \xc9\xe2n\xd9\xf6\x06\x0b8\xf9\xb9\
\xf8\x9b#81z\xf8\x02\x83\x1e\xa2\xf02\x7f\xad\xd7*h\xad9\xf6\x14U\xca\x90\\i\
\xcc~#h\xaa\xb4\x1b['

        s = RandomIO.RandomIO(123456).read(50)

        self.assertEqual(s, string_123456)

    def test_read(self):
        s1 = RandomIO.RandomIO('seed string')

        with self.assertRaises(RuntimeError) as ex:
            s1.read()

        self.assertEqual(
            str(ex.exception),
            'Stream size must be specified if bytes to read is not.')

    def test_dump(self):
        s1 = RandomIO.RandomIO('seed string')
        s2 = RandomIO.RandomIO('seed string')

        file1 = 'file1'
        file2 = 'file2'

        with open(file1, 'wb') as f:
            s1.dump(f, 100)

        with open(file2, 'wb') as f:
            s2.dump(f, 100)

        with open(file1, 'rb') as f:
            contents1 = f.read()

        with open(file2, 'rb') as f:
            contents2 = f.read()

        self.assertEqual(len(contents1), 100)
        self.assertEqual(contents1, contents2)

        os.remove(file1)
        os.remove(file2)

    def test_genfile(self):
        path = RandomIO.RandomIO('seed string').genfile(100)

        with open(path, 'rb') as f:
            contents = f.read()

        self.assertEqual(len(contents), 100)

        os.remove(path)

        dir = 'test_directory/'

        os.mkdir(dir)
        path = RandomIO.RandomIO('seed string').genfile(100, dir)

        (h1, t1) = os.path.split(dir)
        (h2, t2) = os.path.split(path)
        self.assertEqual(h1, h2)

        with open(path, 'rb') as f:
            contents = f.read()

        self.assertEqual(len(contents), 100)

        os.remove(path)
        os.rmdir(dir)

    def test_large(self):
        length = 100000000
        file1 = RandomIO.RandomIO('seed string').genfile(length)
        file2 = RandomIO.RandomIO('seed string').genfile(length)

        with open(file1, 'rb') as f1:
            with open(file1, 'rb') as f2:
                for c in iter(lambda: f1.read(1000), b''):
                    self.assertEqual(c, f2.read(1000))

        os.remove(file1)
        os.remove(file2)

    def test_read_limit(self):
        s1 = RandomIO.RandomIO('seed string', 100)

        s1.seek(90)

        buf1 = s1.read(100)

        self.assertEqual(len(buf1), 10)

    def test_read_zero(self):
        s1 = RandomIO.RandomIO('seed string')

        b = s1.read(0)

        self.assertEqual(len(b), 0)

    def test_seek_beginning(self):
        s1 = RandomIO.RandomIO('seed string')

        buf1 = s1.read(10)

        s1.seek(0)

        buf2 = s1.read(10)

        self.assertEqual(buf1, buf2)

    def test_seek_middle(self):
        s1 = RandomIO.RandomIO('seed string')

        s1.seek(10000)

        buf1 = s1.read(10)

        s1.seek(-10, os.SEEK_CUR)

        buf2 = s1.read(10)

        self.assertEqual(buf1, buf2)

    def test_seek_end_consistency(self):
        s1 = RandomIO.RandomIO('seed string', 100)

        s1.seek(98)

        buf1 = s1.read(10)

        s1.seek(90)

        buf2 = s1.read(10)

        self.assertEqual(buf1, buf2[-2:])

    def test_seek_end(self):
        s1 = RandomIO.RandomIO('seed string', 1000)

        s1.seek(900)

        buf1 = s1.read(10)

        s1.seek(100, os.SEEK_END)

        buf2 = s1.read(10)

        self.assertEqual(buf1, buf2)

    def test_tell_beginning(self):
        s1 = RandomIO.RandomIO('seed string')

        s1.read(100)

        p = s1.tell()

        self.assertEqual(p, 100)

    def test_tell_seek_parity(self):
        s1 = RandomIO.RandomIO('seed string')

        s1.seek(100)

        p = s1.tell()

        self.assertEqual(p, 100)

    def test_seek_end_not_possible(self):
        s1 = RandomIO.RandomIO('seed string')

        with self.assertRaises(RuntimeError) as ex:
            s1.seek(100, os.SEEK_END)

        self.assertEqual(
            str(ex.exception),
            'Cannot seek from end of stream if size is unknown.')

    def test_iotools_txt(self):
        output = 'txt_test.out'
        size = 10485760
        subprocess.call(
            iotools_call + ['pairgen', str(size),
                            '-p', '10', '-o', output])

        with open(output, 'r') as pairsfile:
            for line in pairsfile:
                (hexseed, hash) = line.rstrip().split(' ')
                seed = binascii.unhexlify(hexseed)
                testhash = hashlib.sha256(
                    RandomIO.RandomIO(seed).read(size)).hexdigest()
                self.assertEqual(hash, testhash)
        os.remove(output)

    def test_iotools_redis(self):
        # no redis support for windows, so just pass the test
        if not _platform == 'win32':
            r = redis.StrictRedis(host='localhost', port=6379, db=0)
            output = 'redis_test.out'
            size = 10485760

            subprocess.call(
                iotools_call + ['pairgen', str(size), '-p', '10', '-o', output,
                                '--redis'])
            subprocess.call(
                '{0} {1} | redis-cli --pipe'.format(cat_cmd, output),
                    shell=True)

            for hexseed in r.scan_iter():
                seed = binascii.unhexlify(hexseed)
                testhash = hashlib.sha256(
                    RandomIO.RandomIO(seed).read(size)).hexdigest()
                self.assertEqual(r.get(hexseed).decode('ascii'), testhash)
            os.remove(output)
            r.flushall()

if __name__ == '__main__':
    unittest.main()
<commit_msg>Fixing PEP8 Violation<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# The MIT License (MIT)
#
# Copyright (c) 2014 William T. James for Storj Labs
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import unittest
import os
import redis
import hashlib
import subprocess
import binascii

import RandomIO
from sys import platform as _platform

if _platform.startswith('linux') or _platform == 'darwin':
    cat_cmd = 'cat'
    iotools_call = ['IOTools.py']
elif _platform == 'win32':
    cat_cmd = 'type'
    iotools_call = ['python', os.path.join('../bin', 'IOTools.py')]

class TestRandomIO(unittest.TestCase):

    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_gen(self):
        s = RandomIO.RandomIO()
        b = s.read(100)

        self.assertEqual(len(b), 100)

        self.assertEqual(
            RandomIO.RandomIO(123456).read(100),
            RandomIO.RandomIO(123456).read(100))
        self.assertEqual(RandomIO.RandomIO(b'byte string seed').read(
            100), RandomIO.RandomIO(b'byte string seed').read(100))
        self.assertEqual(RandomIO.RandomIO(1.23456).read(
            100), RandomIO.RandomIO(1.23456).read(100))

    def test_consistent(self):
        s1 = RandomIO.RandomIO('seed string')
        s2 = RandomIO.RandomIO('seed string')
        s3 = RandomIO.RandomIO('seed string')
        s4 = RandomIO.RandomIO('another seed')

        self.assertEqual(s1.read(100), s2.read(100))
        self.assertNotEqual(s3.read(100), s4.read(100))

    def test_crossplatform(self):
        string_seed1 = b'\t\xb0\xef\xd9\x05p\xe1W\x17\x8a9\xc6!;^6\x1d\xadj\
\xb4#n\x1d/\x12+\xe6\xb1\x80\xc86\x06I\xc4!\x8b39\x84E\x1d\x14\xdf\x14e\x12\
\xfa\xf0\r\x1b'

        s = RandomIO.RandomIO('seed1').read(50)

        self.assertEqual(s, string_seed1)

        string_123456 = b'\x18\xb2\xce\x8a \xc9\xe2n\xd9\xf6\x06\x0b8\xf9\xb9\
\xf8\x9b#81z\xf8\x02\x83\x1e\xa2\xf02\x7f\xad\xd7*h\xad9\xf6\x14U\xca\x90\\i\
\xcc~#h\xaa\xb4\x1b['

        s = RandomIO.RandomIO(123456).read(50)

        self.assertEqual(s, string_123456)

    def test_read(self):
        s1 = RandomIO.RandomIO('seed string')

        with self.assertRaises(RuntimeError) as ex:
            s1.read()

        self.assertEqual(
            str(ex.exception),
            'Stream size must be specified if bytes to read is not.')

    def test_dump(self):
        s1 = RandomIO.RandomIO('seed string')
        s2 = RandomIO.RandomIO('seed string')

        file1 = 'file1'
        file2 = 'file2'

        with open(file1, 'wb') as f:
            s1.dump(f, 100)

        with open(file2, 'wb') as f:
            s2.dump(f, 100)

        with open(file1, 'rb') as f:
            contents1 = f.read()

        with open(file2, 'rb') as f:
            contents2 = f.read()

        self.assertEqual(len(contents1), 100)
        self.assertEqual(contents1, contents2)

        os.remove(file1)
        os.remove(file2)

    def test_genfile(self):
        path = RandomIO.RandomIO('seed string').genfile(100)

        with open(path, 'rb') as f:
            contents = f.read()

        self.assertEqual(len(contents), 100)

        os.remove(path)

        dir = 'test_directory/'

        os.mkdir(dir)
        path = RandomIO.RandomIO('seed string').genfile(100, dir)

        (h1, t1) = os.path.split(dir)
        (h2, t2) = os.path.split(path)
        self.assertEqual(h1, h2)

        with open(path, 'rb') as f:
            contents = f.read()

        self.assertEqual(len(contents), 100)

        os.remove(path)
        os.rmdir(dir)

    def test_large(self):
        length = 100000000
        file1 = RandomIO.RandomIO('seed string').genfile(length)
        file2 = RandomIO.RandomIO('seed string').genfile(length)

        with open(file1, 'rb') as f1:
            with open(file1, 'rb') as f2:
                for c in iter(lambda: f1.read(1000), b''):
                    self.assertEqual(c, f2.read(1000))

        os.remove(file1)
        os.remove(file2)

    def test_read_limit(self):
        s1 = RandomIO.RandomIO('seed string', 100)

        s1.seek(90)

        buf1 = s1.read(100)

        self.assertEqual(len(buf1), 10)

    def test_read_zero(self):
        s1 = RandomIO.RandomIO('seed string')

        b = s1.read(0)

        self.assertEqual(len(b), 0)

    def test_seek_beginning(self):
        s1 = RandomIO.RandomIO('seed string')

        buf1 = s1.read(10)

        s1.seek(0)

        buf2 = s1.read(10)

        self.assertEqual(buf1, buf2)

    def test_seek_middle(self):
        s1 = RandomIO.RandomIO('seed string')

        s1.seek(10000)

        buf1 = s1.read(10)

        s1.seek(-10, os.SEEK_CUR)

        buf2 = s1.read(10)

        self.assertEqual(buf1, buf2)

    def test_seek_end_consistency(self):
        s1 = RandomIO.RandomIO('seed string', 100)

        s1.seek(98)

        buf1 = s1.read(10)

        s1.seek(90)

        buf2 = s1.read(10)

        self.assertEqual(buf1, buf2[-2:])

    def test_seek_end(self):
        s1 = RandomIO.RandomIO('seed string', 1000)

        s1.seek(900)

        buf1 = s1.read(10)

        s1.seek(100, os.SEEK_END)

        buf2 = s1.read(10)

        self.assertEqual(buf1, buf2)

    def test_tell_beginning(self):
        s1 = RandomIO.RandomIO('seed string')

        s1.read(100)

        p = s1.tell()

        self.assertEqual(p, 100)

    def test_tell_seek_parity(self):
        s1 = RandomIO.RandomIO('seed string')

        s1.seek(100)

        p = s1.tell()

        self.assertEqual(p, 100)

    def test_seek_end_not_possible(self):
        s1 = RandomIO.RandomIO('seed string')

        with self.assertRaises(RuntimeError) as ex:
            s1.seek(100, os.SEEK_END)

        self.assertEqual(
            str(ex.exception),
            'Cannot seek from end of stream if size is unknown.')

    def test_iotools_txt(self):
        output = 'txt_test.out'
        size = 10485760
        subprocess.call(
            iotools_call + ['pairgen', str(size),
                            '-p', '10', '-o', output])

        with open(output, 'r') as pairsfile:
            for line in pairsfile:
                (hexseed, hash) = line.rstrip().split(' ')
                seed = binascii.unhexlify(hexseed)
                testhash = hashlib.sha256(
                    RandomIO.RandomIO(seed).read(size)).hexdigest()
                self.assertEqual(hash, testhash)
        os.remove(output)

    def test_iotools_redis(self):
        # no redis support for windows, so just pass the test
        if not _platform == 'win32':
            r = redis.StrictRedis(host='localhost', port=6379, db=0)
            output = 'redis_test.out'
            size = 10485760

            subprocess.call(
                iotools_call + ['pairgen', str(size), '-p', '10', '-o', output,
                                '--redis'])
            subprocess.call(
                '{0} {1} | redis-cli --pipe'.format(cat_cmd, output),
                shell=True)

            for hexseed in r.scan_iter():
                seed = binascii.unhexlify(hexseed)
                testhash = hashlib.sha256(
                    RandomIO.RandomIO(seed).read(size)).hexdigest()
                self.assertEqual(r.get(hexseed).decode('ascii'), testhash)
            os.remove(output)
            r.flushall()

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>#!/usr/bin/env python

'''
This module defines the VaspInputSet abstract base class and
a concrete implementation for the Materials Project.  The basic
concept behind an input set is to specify a scheme to generate
a consistent set of Vasp inputs from a structure without further
user intervention. This ensures comparability across runs.
'''

from __future__ import division

__author__="Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Nov 16, 2011"

import os
import abc
import ConfigParser
import json

from pymatgen.io.vaspio import Incar, Poscar, Potcar, Kpoints

class AbstractVaspInputSet(object):
    """
    Abstract base class representing a set of Vasp input parameters.
    The idea is that using a VaspInputSet, a complete set of input files (INPUT, KPOINTS, POSCAR and POTCAR)
    can be generated in an automated fashion for any structure.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def get_poscar(self, structure):
        '''
        Returns Poscar from a structure.
        '''
        return Poscar(structure)

    @abc.abstractmethod
    def get_kpoints(self, structure):
        '''
        Returns Kpoints from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    @abc.abstractmethod
    def get_incar(self, structure):
        '''
        Returns Incar from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    @abc.abstractmethod
    def get_potcar(self, structure):
        '''
        Returns Potcar from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    def get_all_vasp_input(self, structure):
        '''
        Returns all input files as a dict of {filename: file_as_string}

        Arguments:
            structure:
                Structure object

        Returns:
            dict of {filename: file_as_string}, e.g., {'INCAR':'EDIFF=1e-4...'}
        '''
        return {'INCAR':self.get_incar(structure), 'KPOINTS':self.get_kpoints(structure),
                'POSCAR': self.get_poscar(structure), 'POTCAR': self.get_potcar(structure)}

    def write_input(self, structure, output_dir, make_dir_if_not_present = True):
        """
        Writes a set of VASP input to a directory.

        Arguments:
            structure:
                Structure object
            output_dir:
                Directory to output the VASP input files
            make_dir_if_not_present:
                Set to True if you want the directory (and the whole path) to be created if it is not present.
        """
        if make_dir_if_not_present and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        for k,v in self.get_all_vasp_input(structure).items():
            v.write_file(os.path.join(output_dir, k))

class MITVaspInputSet(AbstractVaspInputSet):
    """
    Standard implementation of VaspInputSet utilizing parameters in the MIT High-throughput project.
    The parameters are chosen specifically for a high-throughput project, which means in general smaller
    pseudopotentials were chosen.

    Please refer to A Jain, G. Hautier, C. Moore, S. P. Ong, C. Fischer, T. Mueller, K. A. Persson, G. Ceder (2011).
    A high-throughput infrastructure for density functional theory calculations. Computational Materials Science, 50(8),
    2295-2310. doi:10.1016/j.commatsci.2011.02.023 for more information.
    """

    def __init__(self):
        self.name = "MITMatgen"
        module_dir = os.path.dirname(os.path.abspath(__file__))
        self._config = ConfigParser.SafeConfigParser()
        self._config.optionxform = str
        self._config.readfp(open(os.path.join(module_dir, "VaspInputSets.cfg")))
        self.potcar_settings = dict(self._config.items(self.name + 'POTCAR'))
        self.kpoints_settings = dict(self._config.items(self.name + 'KPOINTS'))
        self.incar_settings = dict(self._config.items(self.name+'INCAR'))
        for key in ['MAGMOM', 'LDAUU', 'LDAUJ', 'LDAUL']:
            self.incar_settings[key] = json.loads(self.incar_settings[key])

    def get_incar(self, structure):
        incar = Incar()
        poscar = Poscar(structure)
        for key, setting in self.incar_settings.items():
            if key == "MAGMOM":
                incar[key] =  [setting.get(site.specie.symbol, 0.6) for site in structure]
            elif key in ['LDAUU', 'LDAUJ', 'LDAUL']:
                incar[key] =  [setting.get(sym, 0) for sym in poscar.site_symbols]
            elif key == "EDIFF":
                incar[key] =  float(setting) * structure.num_sites
            else:
                incar[key] = setting

        has_u = sum(incar['LDAUU']) > 0
        if not has_u:
            for key in incar.keys():
                if key.startswith('LDAU'):
                    del incar[key]

        return incar

    def get_poscar(self, structure):
        return super(MaterialsProjectVaspInputSet, self).get_poscar(structure)

    def get_potcar(self, structure):
        p = self.get_poscar(structure)
        elements = p.site_symbols
        potcar_symbols = []
        for el in elements:
            potcar_symbols.append(self.potcar_settings[el] if el in self.potcar_settings else el)
        return Potcar(potcar_symbols)

    def get_kpoints(self, structure):
        '''
        Writes out a KPOINTS file using the fully automated grid method. Uses Gamma centered meshes
        for hexagonal cells and Monkhorst-Pack grids otherwise.

        Algorithm:
            Uses a simple approach scaling the number of divisions along each
            reciprocal lattice vector proportional to its length. So of N_{grid}
            grid points are desired, then
            N_{grid} = n_{a}*n_{b}*n_{c}
            n_{b} = l_{b}/l_{a} * n_{a}
            n_{c} = l_{c}/l_{a} * n_{a}

            and

            n_{a} = round( ( (l_{a}**2 / (l_{b}*l_{c}) )*N_{grid} )**(1/3) )
            n_{a} = (n_{a} == 0 ? 1 : n_{a})
        '''

        latt = structure.lattice
        lengths = latt.abc
        ngrid = int(self.kpoints_settings['grid_density']) / structure.num_sites

        mult = (ngrid*lengths[0]*lengths[1] * lengths[2]) ** (1/3)

        num_div = [int(round(1/lengths[i] * mult)) for i in xrange(3)]
        #ensure that numDiv[i] > 0

        num_div = [i if i > 0 else 1 for i in num_div]

        angles = latt.angles
        right_angles = [i for i in xrange(3) if abs(angles[i] - 90) < 5]
        hex_angles = [i for i in xrange(3) if abs(angles[i] - 60) < 5 or abs(angles[i] - 120) < 5]

        is_hexagonal = (len(right_angles) == 2 and len(hex_angles) == 1 and lengths[right_angles[0]] == lengths[right_angles[1]])

        style = 'Gamma'
        if not is_hexagonal:
            num_div = [i + i % 2 for i in num_div]
            style = 'Monk'
        comment = "pymatgen generated Materials Project kpoints with grid density = " + self.kpoints_settings['grid_density'] + ' per atom.'
        num_kpts = 0
        return Kpoints(comment, num_kpts, style, [num_div], [0,0,0])

class MaterialsProjectVaspInputSet(MITVaspInputSet):
    """
    Class representing a set of Vasp input parameters.
    The idea is that using a VaspInputSet, a complete set of input files (INPUT, KPOINTS, POSCAR and POTCAR)
    can be generated in an automated fashion for any structure.
    """
    def __init__(self):
        self.name = "MaterialsProject"
        module_dir = os.path.dirname(os.path.abspath(__file__))
        self._config = ConfigParser.SafeConfigParser()
        self._config.optionxform = str
        self._config.readfp(open(os.path.join(module_dir, "VaspInputSets.cfg")))
        self.potcar_settings = dict(self._config.items(self.name + 'POTCAR'))
        self.kpoints_settings = dict(self._config.items(self.name + 'KPOINTS'))
        self.incar_settings = dict(self._config.items(self.name+'INCAR'))
        for key in ['MAGMOM', 'LDAUU', 'LDAUJ', 'LDAUL']:
            self.incar_settings[key] = json.loads(self.incar_settings[key])
<commit_msg>Commented more heavily the two VaspInputSet implementations to clarify the differences between them.<commit_after>#!/usr/bin/env python

'''
This module defines the VaspInputSet abstract base class and
a concrete implementation for the Materials Project.  The basic
concept behind an input set is to specify a scheme to generate
a consistent set of Vasp inputs from a structure without further
user intervention. This ensures comparability across runs.
'''

from __future__ import division

__author__="Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Nov 16, 2011"

import os
import abc
import ConfigParser
import json

from pymatgen.io.vaspio import Incar, Poscar, Potcar, Kpoints

class AbstractVaspInputSet(object):
    """
    Abstract base class representing a set of Vasp input parameters.
    The idea is that using a VaspInputSet, a complete set of input files (INPUT, KPOINTS, POSCAR and POTCAR)
    can be generated in an automated fashion for any structure.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def get_poscar(self, structure):
        '''
        Returns Poscar from a structure.
        '''
        return Poscar(structure)

    @abc.abstractmethod
    def get_kpoints(self, structure):
        '''
        Returns Kpoints from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    @abc.abstractmethod
    def get_incar(self, structure):
        '''
        Returns Incar from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    @abc.abstractmethod
    def get_potcar(self, structure):
        '''
        Returns Potcar from a structure.

        Arguments:
            structure:
                Structure object
        '''
        return

    def get_all_vasp_input(self, structure):
        '''
        Returns all input files as a dict of {filename: file_as_string}

        Arguments:
            structure:
                Structure object

        Returns:
            dict of {filename: file_as_string}, e.g., {'INCAR':'EDIFF=1e-4...'}
        '''
        return {'INCAR':self.get_incar(structure), 'KPOINTS':self.get_kpoints(structure),
                'POSCAR': self.get_poscar(structure), 'POTCAR': self.get_potcar(structure)}

    def write_input(self, structure, output_dir, make_dir_if_not_present = True):
        """
        Writes a set of VASP input to a directory.

        Arguments:
            structure:
                Structure object
            output_dir:
                Directory to output the VASP input files
            make_dir_if_not_present:
                Set to True if you want the directory (and the whole path) to be created if it is not present.
        """
        if make_dir_if_not_present and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        for k,v in self.get_all_vasp_input(structure).items():
            v.write_file(os.path.join(output_dir, k))

class MITVaspInputSet(AbstractVaspInputSet):
    """
    Standard implementation of VaspInputSet utilizing parameters in the MIT High-throughput project.
    The parameters are chosen specifically for a high-throughput project, which means in general smaller
    pseudopotentials were chosen.

    Please refer to A Jain, G. Hautier, C. Moore, S. P. Ong, C. Fischer, T. Mueller, K. A. Persson, G. Ceder (2011).
    A high-throughput infrastructure for density functional theory calculations. Computational Materials Science, 50(8),
    2295-2310. doi:10.1016/j.commatsci.2011.02.023 for more information.
    """

    def __init__(self):
        self.name = "MITMatgen"
        module_dir = os.path.dirname(os.path.abspath(__file__))
        self._config = ConfigParser.SafeConfigParser()
        self._config.optionxform = str
        self._config.readfp(open(os.path.join(module_dir, "VaspInputSets.cfg")))
        self.potcar_settings = dict(self._config.items(self.name + 'POTCAR'))
        self.kpoints_settings = dict(self._config.items(self.name + 'KPOINTS'))
        self.incar_settings = dict(self._config.items(self.name+'INCAR'))
        for key in ['MAGMOM', 'LDAUU', 'LDAUJ', 'LDAUL']:
            self.incar_settings[key] = json.loads(self.incar_settings[key])

    def get_incar(self, structure):
        incar = Incar()
        poscar = Poscar(structure)
        for key, setting in self.incar_settings.items():
            if key == "MAGMOM":
                incar[key] =  [setting.get(site.specie.symbol, 0.6) for site in structure]
            elif key in ['LDAUU', 'LDAUJ', 'LDAUL']:
                incar[key] =  [setting.get(sym, 0) for sym in poscar.site_symbols]
            elif key == "EDIFF":
                incar[key] =  float(setting) * structure.num_sites
            else:
                incar[key] = setting

        has_u = sum(incar['LDAUU']) > 0
        if not has_u:
            for key in incar.keys():
                if key.startswith('LDAU'):
                    del incar[key]

        return incar

    def get_poscar(self, structure):
        return super(MaterialsProjectVaspInputSet, self).get_poscar(structure)

    def get_potcar(self, structure):
        p = self.get_poscar(structure)
        elements = p.site_symbols
        potcar_symbols = []
        for el in elements:
            potcar_symbols.append(self.potcar_settings[el] if el in self.potcar_settings else el)
        return Potcar(potcar_symbols)

    def get_kpoints(self, structure):
        '''
        Writes out a KPOINTS file using the fully automated grid method. Uses Gamma centered meshes
        for hexagonal cells and Monkhorst-Pack grids otherwise.

        Algorithm:
            Uses a simple approach scaling the number of divisions along each
            reciprocal lattice vector proportional to its length.
        '''

        latt = structure.lattice
        lengths = latt.abc
        ngrid = int(self.kpoints_settings['grid_density']) / structure.num_sites

        mult = (ngrid*lengths[0]*lengths[1] * lengths[2]) ** (1/3)

        num_div = [int(round(1/lengths[i] * mult)) for i in xrange(3)]
        #ensure that numDiv[i] > 0

        num_div = [i if i > 0 else 1 for i in num_div]

        angles = latt.angles
        hex_angle_tol = 5 #in degrees
        hex_length_tol = 0.01 #in angstroms
        right_angles = [i for i in xrange(3) if abs(angles[i] - 90) < hex_angle_tol]
        hex_angles = [i for i in xrange(3) if abs(angles[i] - 60) < hex_angle_tol or abs(angles[i] - 120) < hex_angle_tol]

        is_hexagonal = (len(right_angles) == 2 and len(hex_angles) == 1 and abs(lengths[right_angles[0]] == lengths[right_angles[1]]) < hex_length_tol)

        style = 'Gamma'
        if not is_hexagonal:
            num_div = [i + i % 2 for i in num_div]
            style = 'Monk'
        comment = "pymatgen generated Materials Project kpoints with grid density = " + self.kpoints_settings['grid_density'] + ' per atom.'
        num_kpts = 0
        return Kpoints(comment, num_kpts, style, [num_div], [0,0,0])

class MaterialsProjectVaspInputSet(MITVaspInputSet):
    """
    Implementation of VaspInputSet utilizing parameters in the public Materials Project.
    Typically, the psuedopotentials chosen contain more electrons than the MIT parameters,
    and the k-point grid is ~50% more dense.  The LDAUU parameters are also different
    due to the different psps used, which result in different fitted values (even though
    the methodology of fitting is exactly the same as the MIT scheme).
    """
    def __init__(self):
        self.name = "MaterialsProject"
        module_dir = os.path.dirname(os.path.abspath(__file__))
        self._config = ConfigParser.SafeConfigParser()
        self._config.optionxform = str
        self._config.readfp(open(os.path.join(module_dir, "VaspInputSets.cfg")))
        self.potcar_settings = dict(self._config.items(self.name + 'POTCAR'))
        self.kpoints_settings = dict(self._config.items(self.name + 'KPOINTS'))
        self.incar_settings = dict(self._config.items(self.name+'INCAR'))
        for key in ['MAGMOM', 'LDAUU', 'LDAUJ', 'LDAUL']:
            self.incar_settings[key] = json.loads(self.incar_settings[key])
<|endoftext|>


<commit_before>"""
Popular Quotes Database
"""

import pkg_resources

from ..botbase import command
from ..util import SQLiteQuotes

db_file = pkg_resources.resource_filename('pmxbot.popquotes', 'popquotes.sqlite')

def bartletts(lib, nick, qsearch):
	qs = SQLiteQuotes.from_URI('sqlite:' + db_file)
	qs.lib = lib
	qsearch = qsearch.strip()
	if nick == 'pmxbot':
		qt, i, n = qs.quoteLookup()
		if qt:
			if qt.find(':', 0, 15) > -1:
				qt = qt.split(':', 1)[1].strip()
			return qt
	else:
		qt, i, n = qs.quoteLookupWNum(qsearch)
		if qt:
			return u'(%s/%s): %s' % (i, n, qt)

@command('bender', aliases=('bend',), doc='Quote Bender, a la http://en.wikiquote.org/wiki/Futurama')
def bender(client, event, channel, nick, rest):
	return bartletts('bender', nick, rest)

@command('zoidberg', aliases=('zoid',), doc='Quote Zoidberg, a la http://en.wikiquote.org/wiki/Futurama')
def zoidberg(client, event, channel, nick, rest):
	return bartletts('zoid', nick, rest)

@command('simpsons', aliases=('simp',), doc='Quote the Simpsons, a la http://snpp.com/')
def simpsons(client, event, channel, nick, rest):
	return bartletts('simpsons', nick, rest)

@command('hal', aliases=('2001',), doc='HAL 9000')
def hal(client, event, channel, nick, rest):
	return bartletts('hal', nick, rest)

@command('grail', aliases=(), doc='I'' questing baby')
def grail(client, event, channel, nick, rest):
	return bartletts('grail', nick, rest)

@command('anchorman', aliases=(), doc='Quote Anchorman.')
def anchorman(client, event, channel, nick, rest):
	return bartletts('anchorman', nick, rest)

@command('hangover', aliases=(), doc='Quote hangover.')
def hangover(client, event, channel, nick, rest):
	return bartletts('hangover', nick, rest)

@command('R', aliases=('r',), doc='Quote the R mailing list')
def R(client, event, channel, nick, rest):
	return bartletts('R', nick, rest)
<commit_msg>Refactored popquotes commands into a generator function<commit_after>"""
Popular Quotes Database
"""

import itertools

import pkg_resources

from ..botbase import command
from ..util import SQLiteQuotes

db_file = pkg_resources.resource_filename('pmxbot.popquotes', 'popquotes.sqlite')

def bartletts(lib, nick, qsearch):
	qs = SQLiteQuotes.from_URI('sqlite:' + db_file)
	qs.lib = lib
	qsearch = qsearch.strip()
	if nick == 'pmxbot':
		qt, i, n = qs.quoteLookup()
		if qt:
			if qt.find(':', 0, 15) > -1:
				qt = qt.split(':', 1)[1].strip()
			return qt
	else:
		qt, i, n = qs.quoteLookupWNum(qsearch)
		if qt:
			return u'(%s/%s): %s' % (i, n, qt)

# declare all of the popquotes commands
quote_libs = (
	# name, aliases, doc, lib
	('bender', ('bend',), 'Quote Bender, a la http://en.wikiquote.org/wiki/Futurama', 'bender'),
	('zoidberg', ('zoid',), 'Quote Zoidberg, a la http://en.wikiquote.org/wiki/Futurama', 'zoid'),
	('simpsons', ('simp',), 'Quote the Simpsons, a la http://snpp.com/', 'simpsons'),
	('hal', ('2001',), 'HAL 9000', 'hal'),
	('grail', (), 'I questing baby', 'grail'),
	('anchorman', (), 'Quote Anchorman.', 'anchorman'),
	('hangover', (), 'Quote hangover.', 'hangover'),
	('R', ('r',), 'Quote the R mailing list', 'R'),
)

# create the popquotes commands per the declarations above
def make_command(name, aliases, doc, lib):
	cmd_func = lambda client, event, channel, nick, rest: bartletts(lib, nick, rest)
	cmd_func = command(name, aliases=aliases, doc=doc)(cmd_func)
	globals().update({name: cmd_func})
list(itertools.starmap(make_command, quote_libs))
<|endoftext|>