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

import sys, os

print('')
print('DEPRECATION WARNING: script/create_db.py is deprecated in Monocle/Monkey.')
print('Run `{}/alembic upgrade head` instead.'.format(os.path.dirname(sys.executable)))
print('')
<commit_msg>[add]more deprecation messages to scripts/create_db.py and brought back the old Monocle code.<commit_after>import sys, os
from pathlib import Path
from argparse import ArgumentParser

monocle_dir = Path(__file__).resolve().parents[1]
sys.path.append(str(monocle_dir))

from monocle.db import Base, _engine

def parse_args():
    parser = ArgumentParser()
    parser.add_argument(
        '--for-real',
        dest='for_real',
        help='Run the create_db.py script from stock Monocle.',
        action='store_true'
    )
    return parser.parse_args()

def main():
    args = parse_args()

    if args.for_real:
        Base.metadata.create_all(_engine)
        print('Done!')
    else:
        print('')
        print('DEPRECATION WARNING: script/create_db.py is deprecated in Monocle/Monkey.')
        print('')
        print('Run `{}/alembic upgrade head` instead.'.format(os.path.dirname(sys.executable)))
        print('')
        print('Alembic is introduced to ease db migration and to enforce DB consistency.\nDoing changes to your DB schema manually can introduce unecessary human error.')
        print('If you still want to run create_db.py (and you know what you are doing!),\nrerun this script with the flag `--for-real`.')
        print('For help, use the flag `--help`.')
        print('')

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


<commit_before># Copyright 2014 PressLabs SRL
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import re
import os
import errno

from fuse import FuseOSError

from gitfs.utils.decorators.not_in import not_in
from gitfs.utils.decorators.write_operation import write_operation
from gitfs.log import log

from gitfs.events import writers

from .passthrough import PassthroughView, STATS

class CurrentView(PassthroughView):

    def __init__(self, *args, **kwargs):
        super(CurrentView, self).__init__(*args, **kwargs)
        self.dirty = {}

    @write_operation
    @not_in("ignore", check=["old", "new"])
    def rename(self, old, new):
        new = re.sub(self.regex, '', new)
        result = super(CurrentView, self).rename(old, new)

        message = "Rename %s to %s" % (old, new)
        self._stage(**{
            'remove': os.path.split(old)[1],
            'add': new,
            'message': message
        })

        log.debug("CurrentView: Renamed %s to %s", old, new)
        return result

    @write_operation
    @not_in("ignore", check=["target"])
    def symlink(self, name, target):
        result = os.symlink(target, self._full_path(name))

        message = "Create symlink to %s for %s" % (target, name)
        self._stage(add=name, message=message)

        log.debug("CurrentView: Created symlink to %s from %s", name, target)
        return result

    def readlink(self, path):
        log.debug("CurrentView: Read link %s", path)
        return os.readlink(self._full_path(path))

    def getattr(self, path, fh=None):
        full_path = self._full_path(path)
        status = os.lstat(full_path)

        attrs = dict((key, getattr(status, key)) for key in STATS)
        attrs.update({
            'st_uid': self.uid,
            'st_gid': self.gid,
        })

        log.debug("CurrentView: Get attributes %s for %s", str(attrs), path)
        return attrs

    @write_operation
    @not_in("ignore", check=["path"])
    def write(self, path, buf, offset, fh):
        """
            We don't like big big files, so we need to be really carefull
            with them. First we check for offset, then for size. If any of this
            is off limit, raise EFBIG error and delete the file.
        """

        if offset + len(buf) > self.max_size:
            raise FuseOSError(errno.EFBIG)

        result = super(CurrentView, self).write(path, buf, offset, fh)
        self.dirty[fh] = {
            'message': 'Update %s' % path,
        }

        log.debug("CurrentView: Wrote %s to %s", len(buf), path)
        return result

    @write_operation
    @not_in("ignore", check=["path"])
    def mkdir(self, path, mode):
        result = super(CurrentView, self).mkdir(path, mode)

        keep_path = "%s/.keep" % path
        if not os.path.exists(keep_path):
            fh = self.open_for_write(keep_path, os.O_WRONLY | os.O_CREAT)
            super(CurrentView, self).chmod(keep_path, 0644)

            self.dirty[fh] = {
                'message': "Create the %s directory" % path,
            }

            self.release(keep_path, fh)

        log.debug("CurrentView: Created directory %s with mode %s", path, mode)

        return result

    def create(self, path, mode, fi=None):
        fh = self.open_for_write(path, os.O_WRONLY | os.O_CREAT)
        super(CurrentView, self).chmod(path, mode)

        self.dirty[fh] = {
            'message': "Created %s" % path,
        }

        log.debug("CurrentView: Created %s", path)
        return fh

    @write_operation
    @not_in("ignore", check=["path"])
    def chmod(self, path, mode):
        """
        Executes chmod on the file at os level and then it commits the change.
        """
        str_mode = ('%o' % mode)[-4:]
        if str_mode not in ['0755', '0644']:
            raise FuseOSError(errno.EINVAL)

        result = super(CurrentView, self).chmod(path, mode)

        if os.path.isdir(self._full_path(path)):
            return result

        message = 'Chmod to %s on %s' % (str_mode, path)
        self._stage(add=path, message=message)

        log.debug("CurrentView: Change %s mode to %s", path,
                  ('0%o' % mode)[-4:])
        return result

    @write_operation
    @not_in("ignore", check=["path"])
    def fsync(self, path, fdatasync, fh):
        """
        Each time you fsync, a new commit and push are made
        """

        result = super(CurrentView, self).fsync(path, fdatasync, fh)

        message = 'Fsync %s' % path
        self._stage(add=path, message=message)

        log.debug("CurrentView: Fsync %s", path)
        return result

    @write_operation
    @not_in("ignore", check=["path"])
    def open_for_write(self, path, flags):
        global writers
        fh = self.open_for_read(path, flags)
        writers += 1

        log.debug("CurrentView: Open %s for write", path)
        return fh

    def open_for_read(self, path, flags):
        full_path = self._full_path(path)
        log.info("CurrentView: Open %s for read", path)
        return os.open(full_path, flags)

    def open(self, path, flags):
        write_mode = flags & (os.O_WRONLY | os.O_RDWR |
                              os.O_APPEND | os.O_CREAT)
        if write_mode:
            return self.open_for_write(path, flags)
        return self.open_for_read(path, flags)

    def release(self, path, fh):
        """
        Check for path if something was written to. If so, commit and push
        the changed to upstream.
        """

        if fh in self.dirty:
            message = self.dirty[fh]['message']
            del self.dirty[fh]
            global writers
            writers -= 1
            log.debug("CurrentView: Staged %s for commit", path)
            self._stage(add=path, message=message)

        log.debug("CurrentView: Release %s", path)
        return os.close(fh)

    @write_operation
    @not_in("ignore", check=["path"])
    def rmdir(self, path):
        keep_file = os.path.join(path, '.keep')
        # TODO: don`t use the self.unlink method as it will create an unnecessary
        # commit (for the .keep file)
        self.unlink(keep_file)

        result = super(CurrentView, self).rmdir(path)

        message = 'Deleted the %s directory' % path
        log.debug("CurrentView: %s", message)

        return result

    @write_operation
    @not_in("ignore", check=["path"])
    def unlink(self, path):
        result = super(CurrentView, self).unlink(path)

        message = 'Deleted %s' % path
        self._stage(remove=path, message=message)

        log.debug("CurrentView: Deleted %s", path)
        return result

    def _stage(self, message, add=None, remove=None):
        non_empty = False

        add = self._sanitize(add)
        remove = self._sanitize(remove)

        if remove is not None:
            self.repo.index.remove(self._sanitize(remove))
            non_empty = True

        if add is not None:
            self.repo.index.add(self._sanitize(add))
            non_empty = True

        if non_empty:
            self.queue.commit(add=add, remove=remove, message=message)

    def _sanitize(self, path):
        if path is not None and path.startswith("/"):
            path = path[1:]
        return path
<commit_msg>Update comments (rmdir).<commit_after># Copyright 2014 PressLabs SRL
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import re
import os
import errno

from fuse import FuseOSError

from gitfs.utils.decorators.not_in import not_in
from gitfs.utils.decorators.write_operation import write_operation
from gitfs.log import log

from gitfs.events import writers

from .passthrough import PassthroughView, STATS

class CurrentView(PassthroughView):

    def __init__(self, *args, **kwargs):
        super(CurrentView, self).__init__(*args, **kwargs)
        self.dirty = {}

    @write_operation
    @not_in("ignore", check=["old", "new"])
    def rename(self, old, new):
        new = re.sub(self.regex, '', new)
        result = super(CurrentView, self).rename(old, new)

        message = "Rename %s to %s" % (old, new)
        self._stage(**{
            'remove': os.path.split(old)[1],
            'add': new,
            'message': message
        })

        log.debug("CurrentView: Renamed %s to %s", old, new)
        return result

    @write_operation
    @not_in("ignore", check=["target"])
    def symlink(self, name, target):
        result = os.symlink(target, self._full_path(name))

        message = "Create symlink to %s for %s" % (target, name)
        self._stage(add=name, message=message)

        log.debug("CurrentView: Created symlink to %s from %s", name, target)
        return result

    def readlink(self, path):
        log.debug("CurrentView: Read link %s", path)
        return os.readlink(self._full_path(path))

    def getattr(self, path, fh=None):
        full_path = self._full_path(path)
        status = os.lstat(full_path)

        attrs = dict((key, getattr(status, key)) for key in STATS)
        attrs.update({
            'st_uid': self.uid,
            'st_gid': self.gid,
        })

        log.debug("CurrentView: Get attributes %s for %s", str(attrs), path)
        return attrs

    @write_operation
    @not_in("ignore", check=["path"])
    def write(self, path, buf, offset, fh):
        """
            We don't like big big files, so we need to be really carefull
            with them. First we check for offset, then for size. If any of this
            is off limit, raise EFBIG error and delete the file.
        """

        if offset + len(buf) > self.max_size:
            raise FuseOSError(errno.EFBIG)

        result = super(CurrentView, self).write(path, buf, offset, fh)
        self.dirty[fh] = {
            'message': 'Update %s' % path,
        }

        log.debug("CurrentView: Wrote %s to %s", len(buf), path)
        return result

    @write_operation
    @not_in("ignore", check=["path"])
    def mkdir(self, path, mode):
        result = super(CurrentView, self).mkdir(path, mode)

        keep_path = "%s/.keep" % path
        if not os.path.exists(keep_path):
            fh = self.open_for_write(keep_path, os.O_WRONLY | os.O_CREAT)
            super(CurrentView, self).chmod(keep_path, 0644)

            self.dirty[fh] = {
                'message': "Create the %s directory" % path,
            }

            self.release(keep_path, fh)

        log.debug("CurrentView: Created directory %s with mode %s", path, mode)

        return result

    def create(self, path, mode, fi=None):
        fh = self.open_for_write(path, os.O_WRONLY | os.O_CREAT)
        super(CurrentView, self).chmod(path, mode)

        self.dirty[fh] = {
            'message': "Created %s" % path,
        }

        log.debug("CurrentView: Created %s", path)
        return fh

    @write_operation
    @not_in("ignore", check=["path"])
    def chmod(self, path, mode):
        """
        Executes chmod on the file at os level and then it commits the change.
        """
        str_mode = ('%o' % mode)[-4:]
        if str_mode not in ['0755', '0644']:
            raise FuseOSError(errno.EINVAL)

        result = super(CurrentView, self).chmod(path, mode)

        if os.path.isdir(self._full_path(path)):
            return result

        message = 'Chmod to %s on %s' % (str_mode, path)
        self._stage(add=path, message=message)

        log.debug("CurrentView: Change %s mode to %s", path,
                  ('0%o' % mode)[-4:])
        return result

    @write_operation
    @not_in("ignore", check=["path"])
    def fsync(self, path, fdatasync, fh):
        """
        Each time you fsync, a new commit and push are made
        """

        result = super(CurrentView, self).fsync(path, fdatasync, fh)

        message = 'Fsync %s' % path
        self._stage(add=path, message=message)

        log.debug("CurrentView: Fsync %s", path)
        return result

    @write_operation
    @not_in("ignore", check=["path"])
    def open_for_write(self, path, flags):
        global writers
        fh = self.open_for_read(path, flags)
        writers += 1

        log.debug("CurrentView: Open %s for write", path)
        return fh

    def open_for_read(self, path, flags):
        full_path = self._full_path(path)
        log.info("CurrentView: Open %s for read", path)
        return os.open(full_path, flags)

    def open(self, path, flags):
        write_mode = flags & (os.O_WRONLY | os.O_RDWR |
                              os.O_APPEND | os.O_CREAT)
        if write_mode:
            return self.open_for_write(path, flags)
        return self.open_for_read(path, flags)

    def release(self, path, fh):
        """
        Check for path if something was written to. If so, commit and push
        the changed to upstream.
        """

        if fh in self.dirty:
            message = self.dirty[fh]['message']
            del self.dirty[fh]
            global writers
            writers -= 1
            log.debug("CurrentView: Staged %s for commit", path)
            self._stage(add=path, message=message)

        log.debug("CurrentView: Release %s", path)
        return os.close(fh)

    @write_operation
    @not_in("ignore", check=["path"])
    def rmdir(self, path):
        # Unlink the .keep file
        keep_file = os.path.join(path, '.keep')
        self.unlink(keep_file)

        # Delete the actual directory
        result = super(CurrentView, self).rmdir(path)

        message = 'Deleted the %s directory' % path
        log.debug("CurrentView: %s", message)

        return result

    @write_operation
    @not_in("ignore", check=["path"])
    def unlink(self, path):
        result = super(CurrentView, self).unlink(path)

        message = 'Deleted %s' % path
        self._stage(remove=path, message=message)

        log.debug("CurrentView: Deleted %s", path)
        return result

    def _stage(self, message, add=None, remove=None):
        non_empty = False

        add = self._sanitize(add)
        remove = self._sanitize(remove)

        if remove is not None:
            self.repo.index.remove(self._sanitize(remove))
            non_empty = True

        if add is not None:
            self.repo.index.add(self._sanitize(add))
            non_empty = True

        if non_empty:
            self.queue.commit(add=add, remove=remove, message=message)

    def _sanitize(self, path):
        if path is not None and path.startswith("/"):
            path = path[1:]
        return path
<|endoftext|>


<commit_before>import re
import os

from gitfs.filesystems.passthrough import PassthroughFuse, STATS

from .view import View

class CurrentView(PassthroughFuse, View):

    def __init__(self, *args, **kwargs):
        super(CurrentView, self).__init__(*args, **kwargs)
        self.root = self.repo_path
        self.dirty = {}

    def rename(self, old, new):
        new = re.sub(self.regex, '', new)
        super(CurrentView, self).rename(old, new)

    def symlink(self, name, target):
        return os.symlink(target, self._full_path(name))

    def readlink(self, path):
        return os.readlink(self._full_path(path))

    def getattr(self, path, fh=None):
        full_path = self._full_path(path)
        st = os.lstat(full_path)

        attrs = dict((key, getattr(st, key)) for key in STATS)
        attrs.update({
            'st_uid': self.uid,
            'st_gid': self.gid,
        })

        return attrs

    def write(self, path, buf, offset, fh):
        result = super(CurrentView, self).write(path, buf, offset, fh)
        self.dirty[path] = {
            'message': 'Update %s' % path,
            'is_dirty': True
        }
        return result

    def mkdir(self, path, mode):
        result = super(CurrentView, self).mkdir(path, mode)

        path = "%s/.keep" % os.path.split(path)[1]
        if not os.path.exists(path):
            fh = self.create(path, 0644)
            self.release(path, fh)

        return result

    def create(self, path, mode, fi=None):
        result = super(CurrentView, self).create(path, mode, fi)
        self.dirty[path] = {
            'message': "Created %s" % path,
            'is_dirty': True
        }

        return result

    def chmod(self, path, mode):
        """
        Executes chmod on the file at os level and then it commits the change.
        """

        result = super(CurrentView, self).chmod(path, mode)

        str_mode = str(oct(mode))[3:-1]
        message = 'chmod to %s on %s' % (str_mode, path)
        self.commit(path, message)

        return result

    def release(self, path, fh):
        """
        Check for path if something was written to. If so, commit and push
        the changed to upstream.
        """

        # TODO:add them into a queue and commit/push them in another thread
        if path in self.dirty and self.dirty[path]['is_dirty']:
            self.dirty[path]['is_dirty'] = False
            self.commit(path, self.dirty[path]['message'])

        return os.close(fh)

    def commit(self, path, message):
        if path.startswith("/"):
            path = path[1:]

        self.repo.index.add(path)
        self.repo.commit(message, self.author, self.commiter)
        self.repo.push("origin", self.branch)
<commit_msg>Commit on fsync<commit_after>import re
import os

from gitfs.filesystems.passthrough import PassthroughFuse, STATS

from .view import View

class CurrentView(PassthroughFuse, View):

    def __init__(self, *args, **kwargs):
        super(CurrentView, self).__init__(*args, **kwargs)
        self.root = self.repo_path
        self.dirty = {}

    def rename(self, old, new):
        new = re.sub(self.regex, '', new)
        super(CurrentView, self).rename(old, new)

    def symlink(self, name, target):
        return os.symlink(target, self._full_path(name))

    def readlink(self, path):
        return os.readlink(self._full_path(path))

    def getattr(self, path, fh=None):
        full_path = self._full_path(path)
        st = os.lstat(full_path)

        attrs = dict((key, getattr(st, key)) for key in STATS)
        attrs.update({
            'st_uid': self.uid,
            'st_gid': self.gid,
        })

        return attrs

    def write(self, path, buf, offset, fh):
        result = super(CurrentView, self).write(path, buf, offset, fh)
        self.dirty[path] = {
            'message': 'Update %s' % path,
            'is_dirty': True
        }
        return result

    def mkdir(self, path, mode):
        result = super(CurrentView, self).mkdir(path, mode)

        path = "%s/.keep" % os.path.split(path)[1]
        if not os.path.exists(path):
            fh = self.create(path, 0644)
            self.release(path, fh)

        return result

    def create(self, path, mode, fi=None):
        result = super(CurrentView, self).create(path, mode, fi)
        self.dirty[path] = {
            'message': "Created %s" % path,
            'is_dirty': True
        }

        return result

    def chmod(self, path, mode):
        """
        Executes chmod on the file at os level and then it commits the change.
        """

        result = super(CurrentView, self).chmod(path, mode)

        str_mode = str(oct(mode))[3:-1]
        message = 'Chmod to %s on %s' % (str_mode, path)
        self.commit(path, message)

        return result

    def fsync(self, path, fdatasync, fh):
        result = super(CurrentView, self).fsync(path, fdatasync, fh)

        message = 'Write to %s' % path
        self.commit(path, message)

        return result

    def release(self, path, fh):
        """
        Check for path if something was written to. If so, commit and push
        the changed to upstream.
        """

        # TODO:add them into a queue and commit/push them in another thread
        if path in self.dirty and self.dirty[path]['is_dirty']:
            self.dirty[path]['is_dirty'] = False
            self.commit(path, self.dirty[path]['message'])

        return os.close(fh)

    def commit(self, path, message):
        if path.startswith("/"):
            path = path[1:]

        self.repo.index.add(path)
        self.repo.commit(message, self.author, self.commiter)
        self.repo.push("origin", self.branch)
<|endoftext|>


<commit_before>from __future__ import unicode_literals
from html2text import html2text
from bs4 import BeautifulSoup

def parse_raw_document(data):
    doc = BeautifulSoup(data)
    content_node = doc.find(lambda tag: tag.get('id') == 'wb-cont')
    new_doc = BeautifulSoup()
    new_doc.append(content_node.extract())
    return new_doc

def strip_versioning(doc):
    versioning = doc.find(lambda tag: tag.get('class') == ['info'])
    if versioning is not None and ('document' in versioning.text):
        versioning.decompose()
    return doc

def deemphasize_headers(doc):
    for header in doc.findAll(lambda tag: tag.name == 'h1' and tag.get('class') != ['Title-of-Act']):
        header.name = 'h2'
    return doc

def reemphasize_headers(doc):
    for header in doc.find_all(lambda tag: tag.name == 'h6' and tag.get('class') is None):
        header.name = 'h5'
    return doc

def remove_marginal_notes(doc):
    for tag in doc.findAll(lambda tag: tag.name == 'span' and tag.get('class') == ['wb-invisible']):
        tag.decompose()
    return doc

def remove_provision_lists(doc):
    def appropriate_provision(tag):
        if tag.name == 'ul' and tag.get('class') == ['ProvisionList']:
            return tag.li is not None and tag.li.p is not None and tag.li.p.get('class') in (['Section'], ['Subsection'])
        else:
            return False

    for provision_list in [x for x in doc.find_all(appropriate_provision)]:
        nodes_to_insert = []
        for list_entry in [x for x in provision_list.children]:
            if list_entry.name == 'li':
                for list_element in [x for x in list_entry.children]:
                    nodes_to_insert.append(list_element.extract())
            else:
                nodes_to_insert.append(list_entry.extract())

        if provision_list.previous_sibling is not None:
            for node in nodes_to_insert:
                provision_list.previous_sibling.insert_after(node)
        else:
            for i, node in enumerate(nodes_to_insert):
                provision_list.parent.insert(i, node)
        provision_list.decompose()
    return doc

def reformat_definitions(doc):
    definitions = doc.find(lambda tag: tag.name == 'dl' and tag.get('class') == ['Definition'])
    if definitions is None:
        return doc

    def extract_term_nodes():
        for child in definitions.children:
            if child.name == 'dt':
                yield child

    for term_node in extract_term_nodes():
        terms = []
        for para in term_node.children:
            if para.span is None:
                continue
            if para.span.get('class') == ['DefinedTerm']:
                terms.append('**{}**'.format(para.span.text))
            elif para.span.get('class') == ['DefinedTermLink']:
                terms.append('_{}_'.format(para.span.text))

        term_node.clear()
        if len(terms) > 0:
            new_defn = doc.new_tag('p')
            new_defn.string = ' - '.join(terms)
            term_node.append(new_defn)
    return doc

def repoint_links(doc):
    for link in doc.find_all(lambda tag: tag.name == 'a' and tag.get('href', '').startswith('/')):
        link_parts = link['href'].split('/')[1:]
        if len(link_parts) < 3:
            continue

        link_parts.insert(0, 'canada')
        link_parts.insert(3, link_parts[3].split('-')[0])
        link_parts[-1] = '{}.md'.format(link_parts[-1])

        link['href'] = '/' + '/'.join(link_parts)
    return doc

def fix_multipart_headers(doc):
    for header in doc.find_all(lambda tag: tag.name in ('h2', 'h3', 'h4', 'h5', 'h6')):
        elements = []
        current_span = None
        for header_child in [x for x in header.children]:
            if hasattr(header_child, 'name') and header_child.name == 'span':
                if current_span is None:
                    current_span = header_child.extract()
                    elements.append(current_span)
                else:
                    current_span.string = "{} - {}".format(current_span.text, header_child.text)
                    header_child.extract()
            else:
                elements.append(header_child.extract())

        if len(elements) > 0:
            for element in elements:
                header.append(element)
    return doc

def fix_asterisks(doc):
    for sup in doc.find_all(lambda tag: tag.name == 'sup' and tag.text == '*'):
        sup.string = '(*)'
    return doc

rules = [
    strip_versioning,
    deemphasize_headers,
    reemphasize_headers,
    remove_marginal_notes,
    remove_provision_lists,
    reformat_definitions,
    repoint_links,
    fix_multipart_headers,
    fix_asterisks
]

def reformat_document(doc):
    for rule in rules:
        doc = rule(doc)
    return doc

def reformatted_to_markdown(doc):
    return html2text(unicode(str(doc), 'utf8'), bodywidth=0)

def prepare_markdown(doc):
    doc = parse_raw_document(doc)
    doc = reformat_document(doc)
    doc = reformatted_to_markdown(doc)
    return doc
<commit_msg>still better provision lists<commit_after>from __future__ import unicode_literals
from html2text import html2text
from bs4 import BeautifulSoup

def parse_raw_document(data):
    doc = BeautifulSoup(data)
    content_node = doc.find(lambda tag: tag.get('id') == 'wb-cont')
    new_doc = BeautifulSoup()
    new_doc.append(content_node.extract())
    return new_doc

def strip_versioning(doc):
    versioning = doc.find(lambda tag: tag.get('class') == ['info'])
    if versioning is not None and ('document' in versioning.text):
        versioning.decompose()
    return doc

def deemphasize_headers(doc):
    for header in doc.findAll(lambda tag: tag.name == 'h1' and tag.get('class') != ['Title-of-Act']):
        header.name = 'h2'
    return doc

def reemphasize_headers(doc):
    for header in doc.find_all(lambda tag: tag.name == 'h6' and tag.get('class') is None):
        header.name = 'h5'
    return doc

def remove_marginal_notes(doc):
    for tag in doc.findAll(lambda tag: tag.name == 'span' and tag.get('class') == ['wb-invisible']):
        tag.decompose()
    return doc

def remove_provision_lists(doc):
    def appropriate_provision(tag):
        if tag.name == 'ul' and tag.get('class') == ['ProvisionList']:
            return tag.li is not None and \
                   len(tag.li.find_all(lambda tag: tag.name == 'p' and tag.get('class') in (['Section'], ['Subsection']))) > 0
        else:
            return False

    for provision_list in [x for x in doc.find_all(appropriate_provision)]:
        nodes_to_insert = []
        for list_entry in [x for x in provision_list.children]:
            if list_entry.name == 'li':
                for list_element in [x for x in list_entry.children]:
                    nodes_to_insert.append(list_element.extract())
            else:
                nodes_to_insert.append(list_entry.extract())

        if provision_list.previous_sibling is not None:
            for node in nodes_to_insert:
                provision_list.previous_sibling.insert_after(node)
        else:
            for i, node in enumerate(nodes_to_insert):
                provision_list.parent.insert(i, node)
        provision_list.decompose()
    return doc

def reformat_definitions(doc):
    definitions = doc.find(lambda tag: tag.name == 'dl' and tag.get('class') == ['Definition'])
    if definitions is None:
        return doc

    def extract_term_nodes():
        for child in definitions.children:
            if child.name == 'dt':
                yield child

    for term_node in extract_term_nodes():
        terms = []
        for para in term_node.children:
            if para.span is None:
                continue
            if para.span.get('class') == ['DefinedTerm']:
                terms.append('**{}**'.format(para.span.text))
            elif para.span.get('class') == ['DefinedTermLink']:
                terms.append('_{}_'.format(para.span.text))

        term_node.clear()
        if len(terms) > 0:
            new_defn = doc.new_tag('p')
            new_defn.string = ' - '.join(terms)
            term_node.append(new_defn)
    return doc

def repoint_links(doc):
    for link in doc.find_all(lambda tag: tag.name == 'a' and tag.get('href', '').startswith('/')):
        link_parts = link['href'].split('/')[1:]
        if len(link_parts) < 3:
            continue

        link_parts.insert(0, 'canada')
        link_parts.insert(3, link_parts[3].split('-')[0])
        link_parts[-1] = '{}.md'.format(link_parts[-1])

        link['href'] = '/' + '/'.join(link_parts)
    return doc

def fix_multipart_headers(doc):
    for header in doc.find_all(lambda tag: tag.name in ('h2', 'h3', 'h4', 'h5', 'h6')):
        elements = []
        current_span = None
        for header_child in [x for x in header.children]:
            if hasattr(header_child, 'name') and header_child.name == 'span':
                if current_span is None:
                    current_span = header_child.extract()
                    elements.append(current_span)
                else:
                    current_span.string = "{} - {}".format(current_span.text, header_child.text)
                    header_child.extract()
            else:
                elements.append(header_child.extract())

        if len(elements) > 0:
            for element in elements:
                header.append(element)
    return doc

def fix_asterisks(doc):
    for sup in doc.find_all(lambda tag: tag.name == 'sup' and tag.text == '*'):
        sup.string = '(*)'
    return doc

rules = [
    strip_versioning,
    deemphasize_headers,
    reemphasize_headers,
    remove_marginal_notes,
    remove_provision_lists,
    reformat_definitions,
    repoint_links,
    fix_multipart_headers,
    fix_asterisks
]

def reformat_document(doc):
    for rule in rules:
        doc = rule(doc)
    return doc

def reformatted_to_markdown(doc):
    return html2text(unicode(str(doc), 'utf8'), bodywidth=0)

def prepare_markdown(doc):
    doc = parse_raw_document(doc)
    doc = reformat_document(doc)
    doc = reformatted_to_markdown(doc)
    return doc
<|endoftext|>


<commit_before>from __future__ import unicode_literals
from html2text import html2text
from bs4 import BeautifulSoup

def parse_raw_document(data):
    doc = BeautifulSoup(data)
    content_node = doc.find(lambda tag: tag.get('id') == 'wb-cont')
    new_doc = BeautifulSoup()
    new_doc.append(content_node.extract())
    return new_doc

def strip_versioning(doc):
    versioning = doc.find(lambda tag: tag.get('class') == ['info'])
    if versioning is not None and ('document' in versioning.text):
        versioning.decompose()
    return doc

def deemphasize_headers(doc):
    headers = doc.findAll(lambda tag: tag.name == 'h1' and tag.get('class') != ['Title-of-Act'])
    for header in headers:
        header.name = 'h2'
    return doc

def remove_marginal_notes(doc):
    for tag in doc.findAll(lambda tag: tag.name == 'span' and tag.get('class') == ['wb-invisible']):
        tag.decompose()
    return doc

def remove_provision_lists(doc):
    def is_subsection(tag):
        return tag.name == 'p' and tag.get('class') == ['Subsection']

    def appropriate_provision(tag):
        return tag.name == 'ul' and tag.get('class') == ['ProvisionList'] and tag.find_previous_sibling(is_subsection) is None and tag.find_parent(lambda parent: parent.name == 'ul' and parent.get('class') == ['ProvisionList'] and parent.find_previous_sibling(is_subsection) is None) is None

    for provision_list in [x for x in doc.find_all(appropriate_provision)]:
        nodes_to_insert = []
        for list_entry in [x for x in provision_list.children]:
            if list_entry.name == 'li':
                for list_element in [x for x in list_entry.children]:
                    nodes_to_insert.append(list_element.extract())
            else:
                nodes_to_insert.append(list_entry.extract())

        if provision_list.previous_sibling is not None:
            for node in nodes_to_insert:
                provision_list.previous_sibling.insert_after(node)
        else:
            for i, node in enumerate(nodes_to_insert):
                provision_list.parent.insert(i, node)
        provision_list.decompose()
    return doc

def reformat_definitions(doc):
    definitions = doc.find(lambda tag: tag.name == 'dl' and tag.get('class') == ['Definition'])
    if definitions is None:
        return doc

    def extract_term_nodes():
        for child in definitions.children:
            if child.name == 'dt':
                yield child

    for term_node in extract_term_nodes():
        terms = []
        for para in term_node.children:
            if para.span is None:
                continue
            if para.span.get('class') == ['DefinedTerm']:
                terms.append('**{}**'.format(para.span.text))
            elif para.span.get('class') == ['DefinedTermLink']:
                terms.append('_{}_'.format(para.span.text))

        term_node.clear()
        if len(terms) > 0:
            new_defn = doc.new_tag('p')
            new_defn.string = ' - '.join(terms)
            term_node.append(new_defn)
    return doc

def repoint_links(doc):
    for link in doc.find_all(lambda tag: tag.name == 'a' and tag.get('href', '').startswith('/')):
        link_parts = link['href'].split('/')[1:]
        if len(link_parts) < 3:
            continue

        link_parts.insert(0, 'canada')
        link_parts.insert(3, link_parts[3].split('-')[0])
        link_parts[-1] = '{}.md'.format(link_parts[-1])

        link['href'] = '/' + '/'.join(link_parts)
    return doc

def fix_multipart_headers(doc):
    for header in doc.find_all('h2'):
        spans = []
        for header_child in [x for x in header.children]:
            if header_child.name == 'span':
                spans.append(header_child.extract())

        if len(spans) > 0:
            header_text = [span.text for span in spans]
            new_header = doc.new_tag('span')
            new_header.string = ' - '.join(header_text)
            header.append(new_header)
    return doc

def fix_asterisks(doc):
    for sup in doc.find_all(lambda tag: tag.name == 'sup' and tag.text == '*'):
        sup.string = '(*)'
    return doc

rules = [
    strip_versioning,
    deemphasize_headers,
    remove_marginal_notes,
    remove_provision_lists,
    reformat_definitions,
    repoint_links,
    fix_multipart_headers,
    fix_asterisks
]

def reformat_document(doc):
    for rule in rules:
        doc = rule(doc)
    return doc

def reformatted_to_markdown(doc):
    return html2text(unicode(str(doc), 'utf8'), bodywidth=0)

def prepare_markdown(doc):
    doc = parse_raw_document(doc)
    doc = reformat_document(doc)
    doc = reformatted_to_markdown(doc)
    return doc
<commit_msg>handle multipart headers of h3-h6 as well<commit_after>from __future__ import unicode_literals
from html2text import html2text
from bs4 import BeautifulSoup

def parse_raw_document(data):
    doc = BeautifulSoup(data)
    content_node = doc.find(lambda tag: tag.get('id') == 'wb-cont')
    new_doc = BeautifulSoup()
    new_doc.append(content_node.extract())
    return new_doc

def strip_versioning(doc):
    versioning = doc.find(lambda tag: tag.get('class') == ['info'])
    if versioning is not None and ('document' in versioning.text):
        versioning.decompose()
    return doc

def deemphasize_headers(doc):
    headers = doc.findAll(lambda tag: tag.name == 'h1' and tag.get('class') != ['Title-of-Act'])
    for header in headers:
        header.name = 'h2'
    return doc

def remove_marginal_notes(doc):
    for tag in doc.findAll(lambda tag: tag.name == 'span' and tag.get('class') == ['wb-invisible']):
        tag.decompose()
    return doc

def remove_provision_lists(doc):
    def is_subsection(tag):
        return tag.name == 'p' and tag.get('class') == ['Subsection']

    def appropriate_provision(tag):
        return tag.name == 'ul' and tag.get('class') == ['ProvisionList'] and tag.find_previous_sibling(is_subsection) is None and tag.find_parent(lambda parent: parent.name == 'ul' and parent.get('class') == ['ProvisionList'] and parent.find_previous_sibling(is_subsection) is None) is None

    for provision_list in [x for x in doc.find_all(appropriate_provision)]:
        nodes_to_insert = []
        for list_entry in [x for x in provision_list.children]:
            if list_entry.name == 'li':
                for list_element in [x for x in list_entry.children]:
                    nodes_to_insert.append(list_element.extract())
            else:
                nodes_to_insert.append(list_entry.extract())

        if provision_list.previous_sibling is not None:
            for node in nodes_to_insert:
                provision_list.previous_sibling.insert_after(node)
        else:
            for i, node in enumerate(nodes_to_insert):
                provision_list.parent.insert(i, node)
        provision_list.decompose()
    return doc

def reformat_definitions(doc):
    definitions = doc.find(lambda tag: tag.name == 'dl' and tag.get('class') == ['Definition'])
    if definitions is None:
        return doc

    def extract_term_nodes():
        for child in definitions.children:
            if child.name == 'dt':
                yield child

    for term_node in extract_term_nodes():
        terms = []
        for para in term_node.children:
            if para.span is None:
                continue
            if para.span.get('class') == ['DefinedTerm']:
                terms.append('**{}**'.format(para.span.text))
            elif para.span.get('class') == ['DefinedTermLink']:
                terms.append('_{}_'.format(para.span.text))

        term_node.clear()
        if len(terms) > 0:
            new_defn = doc.new_tag('p')
            new_defn.string = ' - '.join(terms)
            term_node.append(new_defn)
    return doc

def repoint_links(doc):
    for link in doc.find_all(lambda tag: tag.name == 'a' and tag.get('href', '').startswith('/')):
        link_parts = link['href'].split('/')[1:]
        if len(link_parts) < 3:
            continue

        link_parts.insert(0, 'canada')
        link_parts.insert(3, link_parts[3].split('-')[0])
        link_parts[-1] = '{}.md'.format(link_parts[-1])

        link['href'] = '/' + '/'.join(link_parts)
    return doc

def fix_multipart_headers(doc):
    for header in doc.find_all(lambda tag: tag.name in ('h2', 'h3', 'h4', 'h5', 'h6')):
        spans = []
        for header_child in [x for x in header.children]:
            if header_child.name == 'span':
                spans.append(header_child.extract())

        if len(spans) > 0:
            header_text = [span.text for span in spans]
            new_header = doc.new_tag('span')
            new_header.string = ' - '.join(header_text)
            header.append(new_header)
    return doc

def fix_asterisks(doc):
    for sup in doc.find_all(lambda tag: tag.name == 'sup' and tag.text == '*'):
        sup.string = '(*)'
    return doc

rules = [
    strip_versioning,
    deemphasize_headers,
    remove_marginal_notes,
    remove_provision_lists,
    reformat_definitions,
    repoint_links,
    fix_multipart_headers,
    fix_asterisks
]

def reformat_document(doc):
    for rule in rules:
        doc = rule(doc)
    return doc

def reformatted_to_markdown(doc):
    return html2text(unicode(str(doc), 'utf8'), bodywidth=0)

def prepare_markdown(doc):
    doc = parse_raw_document(doc)
    doc = reformat_document(doc)
    doc = reformatted_to_markdown(doc)
    return doc
<|endoftext|>


<commit_before>import sys
import operator

from px import px_terminal

from . import px_process

if sys.version_info.major >= 3:
    # For mypy PEP-484 static typing validation
    from typing import List  # NOQA
    from typing import Dict  # NOQA
    from typing import Tuple  # NOQA
    from six import text_type  # NOQA

GROUPS_COUNT = 4

def get_categories(all_processes):
    # type: (List[px_process.PxProcess]) -> List[Tuple[text_type, int]]
    """
    Group processes by pretty names, keeping track of the total rss_kb in each
    group.

    Return the top groups in order plus one "other" which is the sum of the
    rest. The total number of returned groups will be GROUPS_COUNT.
    """

    names_to_kilobytes = {}  # type: Dict[text_type, int]
    for process in all_processes:
        base_kb = 0
        if process.command in names_to_kilobytes:
            base_kb = names_to_kilobytes[process.command]
        else:
            base_kb = 0

        names_to_kilobytes[process.command] = base_kb + process.rss_kb

    sorted_names = sorted(
        names_to_kilobytes.items(), key=operator.itemgetter(1), reverse=True
    )
    other_total_kb = 0
    return_me = []  # type: List[Tuple[text_type, int]]
    for i, name_and_kilobytes in enumerate(sorted_names):
        if i < GROUPS_COUNT - 1:
            return_me.append(name_and_kilobytes)
        else:
            other_total_kb += name_and_kilobytes[1]
    return_me.append(("others", other_total_kb))

    return return_me

def rambar(ram_bar_length, all_processes):
    # type: (int, List[px_process.PxProcess]) -> text_type

    categories = get_categories(all_processes)
    total_kilobytes = 0
    for category in categories:
        total_kilobytes += category[1]

    bar = u""
    for i, category in enumerate(categories):
        name = category[0]
        kilobytes = category[1]

        chars = int(round(ram_bar_length * kilobytes * 1.0 / total_kilobytes))
        if i == len(categories) - 1:
            # Use all remaining chars
            chars = ram_bar_length - px_terminal.visual_length(bar)

        add_to_bar = px_terminal.get_string_of_length(" " + name, chars)
        if i == 0:
            # First red
            add_to_bar = px_terminal.red(add_to_bar)
        elif i == 1:
            # Second yellow
            add_to_bar = px_terminal.yellow(add_to_bar)
        elif i % 2 == 1:
            # Then alternating between normal and inverse video
            add_to_bar = px_terminal.inverse_video(add_to_bar)

        bar += add_to_bar

    return bar
<commit_msg>Better "others" indicator<commit_after>import sys
import operator

from px import px_terminal

from . import px_process

if sys.version_info.major >= 3:
    # For mypy PEP-484 static typing validation
    from typing import List  # NOQA
    from typing import Dict  # NOQA
    from typing import Tuple  # NOQA
    from six import text_type  # NOQA

GROUPS_COUNT = 4

def get_categories(all_processes):
    # type: (List[px_process.PxProcess]) -> List[Tuple[text_type, int]]
    """
    Group processes by pretty names, keeping track of the total rss_kb in each
    group.

    Return the top groups in order plus one "other" which is the sum of the
    rest. The total number of returned groups will be GROUPS_COUNT.
    """

    names_to_kilobytes = {}  # type: Dict[text_type, int]
    for process in all_processes:
        base_kb = 0
        if process.command in names_to_kilobytes:
            base_kb = names_to_kilobytes[process.command]
        else:
            base_kb = 0

        names_to_kilobytes[process.command] = base_kb + process.rss_kb

    sorted_names = sorted(
        names_to_kilobytes.items(), key=operator.itemgetter(1), reverse=True
    )
    other_total_kb = 0
    return_me = []  # type: List[Tuple[text_type, int]]
    for i, name_and_kilobytes in enumerate(sorted_names):
        if i < GROUPS_COUNT - 1:
            return_me.append(name_and_kilobytes)
        else:
            other_total_kb += name_and_kilobytes[1]
    return_me.append(("...", other_total_kb))

    return return_me

def rambar(ram_bar_length, all_processes):
    # type: (int, List[px_process.PxProcess]) -> text_type

    categories = get_categories(all_processes)
    total_kilobytes = 0
    for category in categories:
        total_kilobytes += category[1]

    bar = u""
    for i, category in enumerate(categories):
        name = category[0]
        kilobytes = category[1]

        chars = int(round(ram_bar_length * kilobytes * 1.0 / total_kilobytes))
        if i == len(categories) - 1:
            # Use all remaining chars
            chars = ram_bar_length - px_terminal.visual_length(bar)

        add_to_bar = px_terminal.get_string_of_length(" " + name, chars)
        if i == 0:
            # First red
            add_to_bar = px_terminal.red(add_to_bar)
        elif i == 1:
            # Second yellow
            add_to_bar = px_terminal.yellow(add_to_bar)
        elif i % 2 == 1:
            # Then alternating between normal and inverse video
            add_to_bar = px_terminal.inverse_video(add_to_bar)

        bar += add_to_bar

    return bar
<|endoftext|>


<commit_before>import sys
import operator

from px import px_terminal

from . import px_process

if sys.version_info.major >= 3:
    # For mypy PEP-484 static typing validation
    from typing import List  # NOQA
    from typing import Dict  # NOQA
    from typing import Tuple  # NOQA
    from six import text_type  # NOQA

GROUPS_COUNT = 4

def get_categories(all_processes):
    # type: (List[px_process.PxProcess]) -> List[Tuple[text_type, int]]
    """
    Group processes by pretty names, keeping track of the total rss_kb in each
    group.

    Return the top groups in order plus one "other" which is the sum of the
    rest. The total number of returned groups will be GROUPS_COUNT.
    """

    names_to_kilobytes = {}  # type: Dict[text_type, int]
    for process in all_processes:
        base_kb = 0
        if process.command in names_to_kilobytes:
            base_kb = names_to_kilobytes[process.command]
        else:
            base_kb = 0

        names_to_kilobytes[process.command] = base_kb + process.rss_kb

    sorted_names = sorted(
        names_to_kilobytes.items(), key=operator.itemgetter(1), reverse=True
    )
    other_total_kb = 0
    return_me = []  # type: List[Tuple[text_type, int]]
    for i, name_and_kilobytes in enumerate(sorted_names):
        if i < GROUPS_COUNT - 1:
            return_me.append(name_and_kilobytes)
        else:
            other_total_kb += name_and_kilobytes[1]
    return_me.append(("others", other_total_kb))

    return return_me

def rambar(ram_bar_length, all_processes):
    # type: (int, List[px_process.PxProcess]) -> text_type

    categories = get_categories(all_processes)
    total_kilobytes = 0
    for category in categories:
        total_kilobytes += category[1]

    bar = u""
    for i, category in enumerate(categories):
        name = category[0]
        kilobytes = category[1]

        chars = int(round(ram_bar_length * kilobytes * 1.0 / total_kilobytes))
        if i == len(categories) - 1:
            # Use all remaining chars
            chars = ram_bar_length - px_terminal.visual_length(bar)

        add_to_bar = px_terminal.get_string_of_length(name, chars)
        if i == 0:
            # First red
            add_to_bar = px_terminal.red(add_to_bar)
        elif i == 1:
            # Second yellow
            add_to_bar = px_terminal.yellow(add_to_bar)
        elif i % 2 == 1:
            # Then alternating between normal and inverse video
            add_to_bar = px_terminal.inverse_video(add_to_bar)

        bar += add_to_bar

    return bar
<commit_msg>This looks much better<commit_after>import sys
import operator

from px import px_terminal

from . import px_process

if sys.version_info.major >= 3:
    # For mypy PEP-484 static typing validation
    from typing import List  # NOQA
    from typing import Dict  # NOQA
    from typing import Tuple  # NOQA
    from six import text_type  # NOQA

GROUPS_COUNT = 4

def get_categories(all_processes):
    # type: (List[px_process.PxProcess]) -> List[Tuple[text_type, int]]
    """
    Group processes by pretty names, keeping track of the total rss_kb in each
    group.

    Return the top groups in order plus one "other" which is the sum of the
    rest. The total number of returned groups will be GROUPS_COUNT.
    """

    names_to_kilobytes = {}  # type: Dict[text_type, int]
    for process in all_processes:
        base_kb = 0
        if process.command in names_to_kilobytes:
            base_kb = names_to_kilobytes[process.command]
        else:
            base_kb = 0

        names_to_kilobytes[process.command] = base_kb + process.rss_kb

    sorted_names = sorted(
        names_to_kilobytes.items(), key=operator.itemgetter(1), reverse=True
    )
    other_total_kb = 0
    return_me = []  # type: List[Tuple[text_type, int]]
    for i, name_and_kilobytes in enumerate(sorted_names):
        if i < GROUPS_COUNT - 1:
            return_me.append(name_and_kilobytes)
        else:
            other_total_kb += name_and_kilobytes[1]
    return_me.append(("others", other_total_kb))

    return return_me

def rambar(ram_bar_length, all_processes):
    # type: (int, List[px_process.PxProcess]) -> text_type

    categories = get_categories(all_processes)
    total_kilobytes = 0
    for category in categories:
        total_kilobytes += category[1]

    bar = u""
    for i, category in enumerate(categories):
        name = category[0]
        kilobytes = category[1]

        chars = int(round(ram_bar_length * kilobytes * 1.0 / total_kilobytes))
        if i == len(categories) - 1:
            # Use all remaining chars
            chars = ram_bar_length - px_terminal.visual_length(bar)

        add_to_bar = px_terminal.get_string_of_length(" " + name, chars)
        if i == 0:
            # First red
            add_to_bar = px_terminal.red(add_to_bar)
        elif i == 1:
            # Second yellow
            add_to_bar = px_terminal.yellow(add_to_bar)
        elif i % 2 == 1:
            # Then alternating between normal and inverse video
            add_to_bar = px_terminal.inverse_video(add_to_bar)

        bar += add_to_bar

    return bar
<|endoftext|>


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

########################################################################
#                                                                      #
# This script was written by Thomas Heavey in 2016.                    #
#        theavey@bu.edu     thomasjheavey@gmail.com                    #
#                                                                      #
# Copyright 2016 Thomas J. Heavey IV                                   #
#                                                                      #
# Licensed under the Apache License, Version 2.0 (the "License");      #
# you may not use this file except in compliance with the License.     #
# You may obtain a copy of the License at                              #
#                                                                      #
#    http://www.apache.org/licenses/LICENSE-2.0                        #
#                                                                      #
# Unless required by applicable law or agreed to in writing, software  #
# distributed under the License is distributed on an "AS IS" BASIS,    #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or      #
# implied.                                                             #
# See the License for the specific language governing permissions and  #
# limitations under the License.                                       #
#                                                                      #
########################################################################

import glob
import os
import re

import gromacs.formats
import gromacs.tools
import matplotlib.pyplot as plt
import numpy as np
import MDAnalysis

__version__ = '0.0.2'

# todo define a set of run data as a class
# possibly as a class of mdanalysis universes, but not sure how to work with
# replicas vs. walkers there.

def find_energies():
    """find_energies() is a function that finds all files in the current
    directory that end in a numeral followed by '.edr'. For each of these
    files, it checks if a file named energy(same number).xvg exists, and if
    not, creates it by calling the GROMACS tool gmx energy where input='13'
    corresponds to the total energy at each time."""
    energy_files = glob.glob('*[0-9].edr')
    output_files = []
    for file_name in energy_files:
        output_name = ('energy' + re.search('[0-9]*(?=\.edr)',
                                            file_name).group(0) + '.xvg')
        if not os.path.isfile(output_name):
            gromacs.tools.Energy(f=file_name, o=output_name, input="13")()
        output_files += [output_name]
    output_files.sort()
    output_files.sort(key=len)
    return output_files

def import_energies(output_files):
    """import_energies(file_list) takes a list of .xvg files in the current
    directory, imports their second columns (likely a list of energies at
    consecutive time steps), and returns that as a list of arrays."""
    imported_data = []
    for file_name in output_files:
        xvg_file = gromacs.formats.XVG(filename=file_name)
        imported_data += [xvg_file.array[1]]
    return imported_data

# Run this only if called from the command line
if __name__ == "__main__":
    import argparse
    # todo add argument and code for way to save to a file instead of viewing
    # todo take arguments to change the plotting
    parser = argparse.ArgumentParser(description='A script to plot energy '
                                                 'histograms from a GROMACS '
                                                 'parallel tempering simulation.')
    parser.add_argument('--version', action='version',
                        version='%(prog)s v{}'.format(__version__))
    args = parser.parse_args()

    out_files = find_energies()
    all_data = import_energies(out_files)

    plt.hist(all_data, 50, histtype='stepfilled')

    plt.show()

def combine_energy_files(basename='energy', files=False):
    """combine_energy_files(basename='energy', files=False) is a function that
    combines a set of .xvg files writes a combined .xvg file.
    'basename' is the first part of the name for the xvg files to be combined and
    should differentiate these from any other *.xvg files in the folder.
    Alternatively, the list of files (in the desired order) can be passed in
    with the keyword 'files'."""
    if not files:
        files = glob.glob(basename + '*.xvg')
        files.sort()
        files.sort(key=len)
    data = [gromacs.formats.XVG(filename=files[0]).array[0]]
    data += import_energies(files)
    data = np.array(data)
    gromacs.formats.XVG(array=data).write(filename=basename+'_comb.xvg')

def deconvolve_energies(energyfile='energy_comb.xvg',
                        indexfile='replica_temp.xvg'):
    """deconvolve_energies(energyfile='energy_comb.xvg',
    indexfile='replica_temp.xvg') is a function that takes an xvg files that is
    has n columns of energies likely from a replica exchange simulation where each
    replica remains at a constant temperature (as GROMACS does) and using the n
    data columns of an index xvg file returns an array of the energies where each
    row is now from one 'walker' (continuous coordinates taken by sampling various
    temperatures or other replica conditions)."""
    energies_indexed = gromacs.formats.XVG(filename=energyfile).array
    indices_indexed = gromacs.formats.XVG(filename=indexfile).array.astype(int)
    # todo check for relative start/end points automatically
    length_e = energies_indexed.shape[1]
    length_i = indices_indexed.shape[1]
    ratio = float(length_e) / float(length_i)
    approx_ratio = int(round(ratio))
    if ratio == 1.0:
        deconvolved_energies = energies_indexed[1:][
            indices_indexed[1:],
            np.arange(length_i)]
        e_times = [energies_indexed[0, 0], energies_indexed[0, -1]]
        i_times = [indices_indexed[0, 0], indices_indexed[0, -1]]
    elif ratio > 1:
        if approx_ratio == ratio:
            extra_e = 0
            extra_i = 0
        elif approx_ratio > ratio:
            extra_i = length_i - length_e / approx_ratio
            extra_e = np.mod(length_e, length_i-extra_i)
        elif approx_ratio < ratio:
            extra_e = np.mod(length_e, length_i)
            extra_i = 0
        else:
            raise ImportError('ratio: {}, approx ratio: {}'.format(ratio, approx_ratio))
        deconvolved_energies = energies_indexed[1:, :length_e-extra_e:approx_ratio][
            indices_indexed[1:, :length_i-extra_i],
            np.arange((length_i-extra_i))]
        e_times = [energies_indexed[0, :length_e-extra_e:approx_ratio][0],
                   energies_indexed[0, :length_e-extra_e:approx_ratio][-1]]
        i_times = [indices_indexed[0, :length_i-extra_i][0],
                   indices_indexed[0, :length_i-extra_i][-1]]
    elif ratio < 1:
        print('likely undersampling energies because energy / indices ratio is '
              '{}'.format(ratio))
        ratio = 1 / ratio
        approx_ratio = int(round(ratio))
        if approx_ratio == ratio:
            extra_e = 0
            extra_i = 0
        # Not so sure about this...
        elif approx_ratio > ratio:
            extra_e = length_e - length_i / approx_ratio
            extra_i = np.mod(length_i, length_e-extra_e)
        elif approx_ratio < ratio:
            extra_i = np.mod(length_i, length_e)
            extra_e = 0
        else:
            raise ImportError('ratio: {}, approx ratio: {}'.format(ratio, approx_ratio))
        deconvolved_energies = energies_indexed[1:, :length_e-extra_e][
            indices_indexed[1:, :length_i-extra_i:approx_ratio],
            np.arange((length_i-extra_i)/approx_ratio)]
        e_times = [energies_indexed[0, :length_e-extra_e][0],
                   energies_indexed[0, :length_e-extra_e][-1]]
        i_times = [indices_indexed[0, :length_i-extra_i:approx_ratio][0],
                   indices_indexed[0, :length_i-extra_i:approx_ratio][-1]]
    else:
        print('length of energy file is {}'.format(length_e))
        print('length of index file is {}'.format(length_i))
        raise ImportError('Not sure how to handle those values')
    # todo maybe error check is a more pythonic manner with a try/except loop
    # if length_e != length_i:
    #     print("length of energies, {}, != length of indices, "
    #           "{}!".format(length_e, length_i))
    #     raise IndexError('lengths not equals')
    print('energies start: {}; end: {}'.format(e_times[0], e_times[1]))
    print('indices start: {}; end: {}'.format(i_times[0], i_times[1]))
    print('These values should be about the same if this is working properly')
    return deconvolved_energies

def plot_array(array, index_offset=0, num_replicas=False, n_rows=False, n_cols=False):
    """plot_array(array, index_offset=0, num_replicas=16, n_rows=False, n_cols=False)
    will put each column of array in a different axes of a figure and then return
    the figure."""
    if not num_replicas:
        num_replicas = array.shape[0] - index_offset
    from math import sqrt, ceil
    if n_rows == n_cols == False:
        n_rows = int(ceil(sqrt(float(num_replicas))))
        n_cols = n_rows
    fig, axes = plt.subplots(n_rows, n_cols, sharex=True, sharey=True)
    for i in range(num_replicas):
        ax = axes.flat[i]
        ax.plot(array[i+index_offset])
    return fig

def hist_array(array, index_offset=0, num_replicas=False, n_rows=False, n_cols=False,
               n_bins=10):
    """hist_array(array, index_offset=0, num_replicas=16, n_rows=False, n_cols=False,
    n_bins=10) will put each column of array in a different axes of a figure and
    then return the figure."""
    if not num_replicas:
        num_replicas = array.shape[0] - index_offset
    from math import sqrt, ceil
    if n_rows == n_cols == False:
        n_rows = int(ceil(sqrt(float(num_replicas))))
        n_cols = n_rows
    fig, axes = plt.subplots(n_rows, n_cols, sharex=True, sharey=True)
    for i in range(num_replicas):
        ax = axes.flat[i]
        ax.hist(array[i+index_offset], n_bins)
    return fig

def solute_trr(trr_base_name='npt_PT_out', tpr_base_name='TOPO/npt',
               output_base_name='solute', index='index.ndx', demux=True):
    """solute_trr takes file base names as input, creates a separate trr file for each
    trajectory that only includes the solutes, and then returns a list of the names of
    the created files.
    If demux=True, it will first use trjcat to deconvolve the walker trajectories (to
    get continuous coordinate files as opposed to continuous temperature).
    This uses gromacswrapper to call trjconv and possibly trjcat."""
    trr_files = glob.glob(trr_base_name + '*.trr')
    trr_files.sort()
    trr_files.sort(key=len)
    tpr_files = glob.glob(tpr_base_name + '*.tpr')
    tpr_files.sort()
    tpr_files.sort(key=len)
    output_files = []
    if demux:
        d_trr_base_name = 'deconv' + trr_base_name
        prev_deconv_files = glob.glob(d_trr_base_name+'*.trr')
        if len(prev_deconv_files) == len(trr_files):
            print('Likely already deconvolved trajectories, skipping that step')
        else:
            gromacs.tools.Trjcat(f=trr_files, o='demuxed.trr', n='index.ndx',
                                 demux='replica_index.xvg', input='CHR')()
            trr_files = glob.glob('*demuxed.trr')
            trr_files.sort()
            trr_files.sort(key=len)
            for (i, trr_file) in enumerate(trr_files):
                number = trr_file.split('_')[0]
                new_name = d_trr_base_name + number + '.trr'
                os.rename(trr_file, new_name)
                trr_files[i] = new_name
        trr_base_name = d_trr_base_name
    if len(trr_files) != len(tpr_files):
        raise IndexError('Number of trr and tpr files not equal: '
                         '{} and {}'.format(len(trr_files), len(tpr_files)))
    for (i, trr_name) in enumerate(trr_files):
        number_match = re.search('(?:'+trr_base_name+')(\d+)(?:\.trr)', trr_name)
        number = number_match.group(1)
        out_file = output_base_name + number + '.trr'
        output_files.append(out_file)
        gromacs.tools.Trjconv(s=tpr_files[i], pbc='mol', f=trr_name, o=out_file,
                              n=index, center=True, input=('CHR', 'CHR'))()
    return output_files

def radii_of_gyration(basename='solute', atom_selection=False, resname='TAD'):
    """A function to find the radius of gyration for all timesteps for a set of
    REMD trajectories.
    The basename is the name before the numbers for the trajectory files.
    resname is the name of the desired residue to be selected.
    A more general atom_selection can be given, in which case resname will
    be ignored.
    Assumes all trajectories are the same length.
    Returns the values as a numpy array."""
    u_solutes = []
    files = glob.glob(basename+'*.trr')
    num_files = len(files)
    for file_name in files:
        u_solutes.append(MDAnalysis.Universe('geom-solutes.gro', file_name))
    rgs = np.zeros((num_files, len(u_solutes[0].trajectory)))
    if not atom_selection:
        selection = 'resname ' + resname
    else:
        selection = atom_selection
    for (i, u) in enumerate(u_solutes):
        tad = u.select_atoms(selection)
        for fr in u.trajectory:
            rgs[i, fr.frame] = tad.radius_of_gyration()
    return rgs
<commit_msg>added basic index alignment check to avoid printed message to deconvolve_energies<commit_after>#! /usr/bin/env python

########################################################################
#                                                                      #
# This script was written by Thomas Heavey in 2016.                    #
#        theavey@bu.edu     thomasjheavey@gmail.com                    #
#                                                                      #
# Copyright 2016 Thomas J. Heavey IV                                   #
#                                                                      #
# Licensed under the Apache License, Version 2.0 (the "License");      #
# you may not use this file except in compliance with the License.     #
# You may obtain a copy of the License at                              #
#                                                                      #
#    http://www.apache.org/licenses/LICENSE-2.0                        #
#                                                                      #
# Unless required by applicable law or agreed to in writing, software  #
# distributed under the License is distributed on an "AS IS" BASIS,    #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or      #
# implied.                                                             #
# See the License for the specific language governing permissions and  #
# limitations under the License.                                       #
#                                                                      #
########################################################################

import glob
import os
import re

import gromacs.formats
import gromacs.tools
import matplotlib.pyplot as plt
import numpy as np
import MDAnalysis

__version__ = '0.0.2'

# todo define a set of run data as a class
# possibly as a class of mdanalysis universes, but not sure how to work with
# replicas vs. walkers there.

def find_energies():
    """find_energies() is a function that finds all files in the current
    directory that end in a numeral followed by '.edr'. For each of these
    files, it checks if a file named energy(same number).xvg exists, and if
    not, creates it by calling the GROMACS tool gmx energy where input='13'
    corresponds to the total energy at each time."""
    energy_files = glob.glob('*[0-9].edr')
    output_files = []
    for file_name in energy_files:
        output_name = ('energy' + re.search('[0-9]*(?=\.edr)',
                                            file_name).group(0) + '.xvg')
        if not os.path.isfile(output_name):
            gromacs.tools.Energy(f=file_name, o=output_name, input="13")()
        output_files += [output_name]
    output_files.sort()
    output_files.sort(key=len)
    return output_files

def import_energies(output_files):
    """import_energies(file_list) takes a list of .xvg files in the current
    directory, imports their second columns (likely a list of energies at
    consecutive time steps), and returns that as a list of arrays."""
    imported_data = []
    for file_name in output_files:
        xvg_file = gromacs.formats.XVG(filename=file_name)
        imported_data += [xvg_file.array[1]]
    return imported_data

# Run this only if called from the command line
if __name__ == "__main__":
    import argparse
    # todo add argument and code for way to save to a file instead of viewing
    # todo take arguments to change the plotting
    parser = argparse.ArgumentParser(description='A script to plot energy '
                                                 'histograms from a GROMACS '
                                                 'parallel tempering simulation.')
    parser.add_argument('--version', action='version',
                        version='%(prog)s v{}'.format(__version__))
    args = parser.parse_args()

    out_files = find_energies()
    all_data = import_energies(out_files)

    plt.hist(all_data, 50, histtype='stepfilled')

    plt.show()

def combine_energy_files(basename='energy', files=False):
    """combine_energy_files(basename='energy', files=False) is a function that
    combines a set of .xvg files writes a combined .xvg file.
    'basename' is the first part of the name for the xvg files to be combined and
    should differentiate these from any other *.xvg files in the folder.
    Alternatively, the list of files (in the desired order) can be passed in
    with the keyword 'files'."""
    if not files:
        files = glob.glob(basename + '*.xvg')
        files.sort()
        files.sort(key=len)
    data = [gromacs.formats.XVG(filename=files[0]).array[0]]
    data += import_energies(files)
    data = np.array(data)
    gromacs.formats.XVG(array=data).write(filename=basename+'_comb.xvg')

def deconvolve_energies(energyfile='energy_comb.xvg',
                        indexfile='replica_temp.xvg'):
    """deconvolve_energies(energyfile='energy_comb.xvg',
    indexfile='replica_temp.xvg') is a function that takes an xvg files that is
    has n columns of energies likely from a replica exchange simulation where each
    replica remains at a constant temperature (as GROMACS does) and using the n
    data columns of an index xvg file returns an array of the energies where each
    row is now from one 'walker' (continuous coordinates taken by sampling various
    temperatures or other replica conditions)."""
    energies_indexed = gromacs.formats.XVG(filename=energyfile).array
    indices_indexed = gromacs.formats.XVG(filename=indexfile).array.astype(int)
    # todo check for relative start/end points automatically
    length_e = energies_indexed.shape[1]
    length_i = indices_indexed.shape[1]
    ratio = float(length_e) / float(length_i)
    approx_ratio = int(round(ratio))
    if ratio == 1.0:
        deconvolved_energies = energies_indexed[1:][
            indices_indexed[1:],
            np.arange(length_i)]
        e_times = [energies_indexed[0, 0], energies_indexed[0, -1]]
        i_times = [indices_indexed[0, 0], indices_indexed[0, -1]]
    elif ratio > 1:
        if approx_ratio == ratio:
            extra_e = 0
            extra_i = 0
        elif approx_ratio > ratio:
            extra_i = length_i - length_e / approx_ratio
            extra_e = np.mod(length_e, length_i-extra_i)
        elif approx_ratio < ratio:
            extra_e = np.mod(length_e, length_i)
            extra_i = 0
        else:
            raise ImportError('ratio: {}, approx ratio: {}'.format(ratio, approx_ratio))
        deconvolved_energies = energies_indexed[1:, :length_e-extra_e:approx_ratio][
            indices_indexed[1:, :length_i-extra_i],
            np.arange((length_i-extra_i))]
        e_times = [energies_indexed[0, :length_e-extra_e:approx_ratio][0],
                   energies_indexed[0, :length_e-extra_e:approx_ratio][-1]]
        i_times = [indices_indexed[0, :length_i-extra_i][0],
                   indices_indexed[0, :length_i-extra_i][-1]]
    elif ratio < 1:
        print('likely undersampling energies because energy / indices ratio is '
              '{}'.format(ratio))
        ratio = 1 / ratio
        approx_ratio = int(round(ratio))
        if approx_ratio == ratio:
            extra_e = 0
            extra_i = 0
        # Not so sure about this...
        elif approx_ratio > ratio:
            extra_e = length_e - length_i / approx_ratio
            extra_i = np.mod(length_i, length_e-extra_e)
        elif approx_ratio < ratio:
            extra_i = np.mod(length_i, length_e)
            extra_e = 0
        else:
            raise ImportError('ratio: {}, approx ratio: {}'.format(ratio, approx_ratio))
        deconvolved_energies = energies_indexed[1:, :length_e-extra_e][
            indices_indexed[1:, :length_i-extra_i:approx_ratio],
            np.arange((length_i-extra_i)/approx_ratio)]
        e_times = [energies_indexed[0, :length_e-extra_e][0],
                   energies_indexed[0, :length_e-extra_e][-1]]
        i_times = [indices_indexed[0, :length_i-extra_i:approx_ratio][0],
                   indices_indexed[0, :length_i-extra_i:approx_ratio][-1]]
    else:
        print('length of energy file is {}'.format(length_e))
        print('length of index file is {}'.format(length_i))
        raise ImportError('Not sure how to handle those values')
    # todo maybe error check is a more pythonic manner with a try/except loop
    # if length_e != length_i:
    #     print("length of energies, {}, != length of indices, "
    #           "{}!".format(length_e, length_i))
    #     raise IndexError('lengths not equals')
    if not (float(e_times[0]) == float(i_times[0]) and
            float(e_times[1]) == float(i_times[1])):
        print('energies start: {}; end: {}'.format(e_times[0], e_times[1]))
        print('indices start: {}; end: {}'.format(i_times[0], i_times[1]))
        print('These values should be about the same if this is working properly')
    return deconvolved_energies

def plot_array(array, index_offset=0, num_replicas=False, n_rows=False, n_cols=False):
    """plot_array(array, index_offset=0, num_replicas=16, n_rows=False, n_cols=False)
    will put each column of array in a different axes of a figure and then return
    the figure."""
    if not num_replicas:
        num_replicas = array.shape[0] - index_offset
    from math import sqrt, ceil
    if n_rows == n_cols == False:
        n_rows = int(ceil(sqrt(float(num_replicas))))
        n_cols = n_rows
    fig, axes = plt.subplots(n_rows, n_cols, sharex=True, sharey=True)
    for i in range(num_replicas):
        ax = axes.flat[i]
        ax.plot(array[i+index_offset])
    return fig

def hist_array(array, index_offset=0, num_replicas=False, n_rows=False, n_cols=False,
               n_bins=10):
    """hist_array(array, index_offset=0, num_replicas=16, n_rows=False, n_cols=False,
    n_bins=10) will put each column of array in a different axes of a figure and
    then return the figure."""
    if not num_replicas:
        num_replicas = array.shape[0] - index_offset
    from math import sqrt, ceil
    if n_rows == n_cols == False:
        n_rows = int(ceil(sqrt(float(num_replicas))))
        n_cols = n_rows
    fig, axes = plt.subplots(n_rows, n_cols, sharex=True, sharey=True)
    for i in range(num_replicas):
        ax = axes.flat[i]
        ax.hist(array[i+index_offset], n_bins)
    return fig

def solute_trr(trr_base_name='npt_PT_out', tpr_base_name='TOPO/npt',
               output_base_name='solute', index='index.ndx', demux=True):
    """solute_trr takes file base names as input, creates a separate trr file for each
    trajectory that only includes the solutes, and then returns a list of the names of
    the created files.
    If demux=True, it will first use trjcat to deconvolve the walker trajectories (to
    get continuous coordinate files as opposed to continuous temperature).
    This uses gromacswrapper to call trjconv and possibly trjcat."""
    trr_files = glob.glob(trr_base_name + '*.trr')
    trr_files.sort()
    trr_files.sort(key=len)
    tpr_files = glob.glob(tpr_base_name + '*.tpr')
    tpr_files.sort()
    tpr_files.sort(key=len)
    output_files = []
    if demux:
        d_trr_base_name = 'deconv' + trr_base_name
        prev_deconv_files = glob.glob(d_trr_base_name+'*.trr')
        if len(prev_deconv_files) == len(trr_files):
            print('Likely already deconvolved trajectories, skipping that step')
        else:
            gromacs.tools.Trjcat(f=trr_files, o='demuxed.trr', n='index.ndx',
                                 demux='replica_index.xvg', input='CHR')()
            trr_files = glob.glob('*demuxed.trr')
            trr_files.sort()
            trr_files.sort(key=len)
            for (i, trr_file) in enumerate(trr_files):
                number = trr_file.split('_')[0]
                new_name = d_trr_base_name + number + '.trr'
                os.rename(trr_file, new_name)
                trr_files[i] = new_name
        trr_base_name = d_trr_base_name
    if len(trr_files) != len(tpr_files):
        raise IndexError('Number of trr and tpr files not equal: '
                         '{} and {}'.format(len(trr_files), len(tpr_files)))
    for (i, trr_name) in enumerate(trr_files):
        number_match = re.search('(?:'+trr_base_name+')(\d+)(?:\.trr)', trr_name)
        number = number_match.group(1)
        out_file = output_base_name + number + '.trr'
        output_files.append(out_file)
        gromacs.tools.Trjconv(s=tpr_files[i], pbc='mol', f=trr_name, o=out_file,
                              n=index, center=True, input=('CHR', 'CHR'))()
    return output_files

def radii_of_gyration(basename='solute', atom_selection=False, resname='TAD'):
    """A function to find the radius of gyration for all timesteps for a set of
    REMD trajectories.
    The basename is the name before the numbers for the trajectory files.
    resname is the name of the desired residue to be selected.
    A more general atom_selection can be given, in which case resname will
    be ignored.
    Assumes all trajectories are the same length.
    Returns the values as a numpy array."""
    u_solutes = []
    files = glob.glob(basename+'*.trr')
    num_files = len(files)
    for file_name in files:
        u_solutes.append(MDAnalysis.Universe('geom-solutes.gro', file_name))
    rgs = np.zeros((num_files, len(u_solutes[0].trajectory)))
    if not atom_selection:
        selection = 'resname ' + resname
    else:
        selection = atom_selection
    for (i, u) in enumerate(u_solutes):
        tad = u.select_atoms(selection)
        for fr in u.trajectory:
            rgs[i, fr.frame] = tad.radius_of_gyration()
    return rgs
<|endoftext|>