<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright 2018 <NAME>
#
# 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.

"""
Modifies the file to comply (as much as possible) with the format specifications.

This includes:
    - Adding TER statements after chain breaks/changes
    - Truncating/Padding all lines to 80 characters
    - Adds END statement at the end of the file

This will remove all original TER/END statements from the file.

Usage:
    python pdb_tidy.py <pdb file>

Example:
    python pdb_tidy.py 1CTF.pdb

This program is part of the `pdb-tools` suite of utilities and should not be
distributed isolatedly. The `pdb-tools` were created to quickly manipulate PDB
files using the terminal, and can be used sequentially, with one tool streaming
data to another. They are based on old FORTRAN77 code that was taking too much
effort to maintain and compile. RIP.
"""

import os
import sys

__author__ = "<NAME>"
__email__ = "<EMAIL>"


def check_input(args):
    """Checks whether to read from stdin/file and validates user input/options.
    """

    # Defaults
    fh = sys.stdin  # file handle

    if not len(args):
        # Reading from pipe with default option
        if sys.stdin.isatty():
            sys.stderr.write(__doc__)
            sys.exit(1)

    elif len(args) == 1:
        if not os.path.isfile(args[0]):
            emsg = 'ERROR!! File not found or not readable: \'{}\'\n'
            sys.stderr.write(emsg.format(args[0]))
            sys.stderr.write(__doc__)
            sys.exit(1)

        fh = open(args[0], 'r')

    else:  # Whatever ...
        sys.stderr.write(__doc__)
        sys.exit(1)

    return fh


def tidy_pdbfile(fhandle):
    """Adds TER/END statements and pads all lines to 80 characters.
    """

    # TER     606      LEU A  75
    fmt_TER = "TER   {:>5d}      {:3s} {:1s}{:>4s}{:1s}" + " " * 53 + "\n"
    # CONECT 1179  746 1184 1195 1203
    fmt_CONECT = "CONECT{:>5d}{:>5d}{:>5d}{:>5d}{:>5d}" + " " * 49 + "\n"
    char_ranges = (slice(6, 11), slice(11, 16),
                   slice(16, 21), slice(21, 26), slice(26, 31))

    records = ('ATOM', 'HETATM')
    ignored = ('TER', 'END')
    # Iterate up to the first ATOM/HETATM line
    prev_line = None
    for line in fhandle:
        if line.startswith(ignored):
            continue

        # Check line length
        line_length = len(line)
        if line_length < 80:
            num_pads = 80 - line_length + 1
            line = line + ' ' * num_pads

        yield line[:80]

        if line.startswith(records):
            prev_line = line
            break

    # Now go through all the remaining lines
    serial_equiv = {}  # store for conect statements
    coord_section = False
    serial_offset = 0  # To offset after adding TER records
    for line in fhandle:
        if line.startswith(ignored):
            continue

        elif line.startswith(records):
            coord_section = True
            is_gap = (int(line[22:26]) - int(prev_line[22:26])) > 1
            if line[21] != prev_line[21] or is_gap:
                serial = int(prev_line[6:11]) + 1
                rname = prev_line[17:20]
                chain = prev_line[21]
                resid = prev_line[22:26]
                icode = prev_line[26]

                serial_offset += 1

                ter_line = fmt_TER.format(serial, rname, chain, resid, icode)
                yield ter_line

            ori_serial = int(prev_line[6:11])
            serial = ori_serial + serial_offset

            serial_equiv[ori_serial] = serial

            line = line[:6] + str(serial).rjust(5) + line[11:]
            prev_line = line
            serial_offset = 1  # reset otherwise we have +2 ex.

        elif line.startswith('ANISOU'):
            # Fix serial based on previous atom
            serial = prev_line[6:11]
            line = line[:6] + str(serial).rjust(5) + line[11:]

        elif line.startswith('CONECT'):
            # 6:11, 11:16, 16:21, 21:26, 26:31
            serials = [line[cr] for cr in char_ranges]
            # If not found, return default
            new_serials = [serial_equiv.get(s, s) for s in serials]
            conect_line = fmt_CONECT.format(*new_serials)
            yield conect_line
            continue

        else:
            if coord_section:
                coord_section = False
                # Add last TER statement
                serial = int(prev_line[6:11]) + 1
                rname = prev_line[17:20]
                chain = prev_line[21]
                resid = prev_line[22:26]
                icode = prev_line[26]

                ter_line = fmt_TER.format(serial, rname, chain, resid, icode)
                yield ter_line

        # Check line length
        line_length = len(line)
        if line_length < 80:
            num_pads = 80 - line_length + 1
            line = line + ' ' * num_pads

        yield line

    # Add END statement
    yield 'END' + ' ' * 77 + '\n'


def main():
    # Check Input
    pdbfh = check_input(sys.argv[1:])<commit_msg>improved line formatting and corrected newline in HEADER<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright 2018 <NAME>
#
# 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.

"""
Modifies the file to comply (as much as possible) with the format specifications.

This includes:
    - Adding TER statements after chain breaks/changes
    - Truncating/Padding all lines to 80 characters
    - Adds END statement at the end of the file

This will remove all original TER/END statements from the file.

Usage:
    python pdb_tidy.py <pdb file>

Example:
    python pdb_tidy.py 1CTF.pdb

This program is part of the `pdb-tools` suite of utilities and should not be
distributed isolatedly. The `pdb-tools` were created to quickly manipulate PDB
files using the terminal, and can be used sequentially, with one tool streaming
data to another. They are based on old FORTRAN77 code that was taking too much
effort to maintain and compile. RIP.
"""

import os
import sys

__author__ = "<NAME>"
__email__ = "<EMAIL>"


def check_input(args):
    """Checks whether to read from stdin/file and validates user input/options.
    """

    # Defaults
    fh = sys.stdin  # file handle

    if not len(args):
        # Reading from pipe with default option
        if sys.stdin.isatty():
            sys.stderr.write(__doc__)
            sys.exit(1)

    elif len(args) == 1:
        if not os.path.isfile(args[0]):
            emsg = 'ERROR!! File not found or not readable: \'{}\'\n'
            sys.stderr.write(emsg.format(args[0]))
            sys.stderr.write(__doc__)
            sys.exit(1)

        fh = open(args[0], 'r')

    else:  # Whatever ...
        sys.stderr.write(__doc__)
        sys.exit(1)

    return fh


def tidy_pdbfile(fhandle):
    """Adds TER/END statements and pads all lines to 80 characters.
    """

    # TER     606      LEU A  75
    fmt_TER = "TER   {:>5d}      {:3s} {:1s}{:>4s}{:1s}" + " " * 53 + "\n"
    # CONECT 1179  746 1184 1195 1203
    fmt_CONECT = "CONECT{:>5d}{:>5d}{:>5d}{:>5d}{:>5d}" + " " * 49 + "\n"
    char_ranges = (slice(6, 11), slice(11, 16),
                   slice(16, 21), slice(21, 26), slice(26, 31))

    records = ('ATOM', 'HETATM')
    ignored = ('TER', 'END')
    # Iterate up to the first ATOM/HETATM line
    prev_line = None
    for line in fhandle:
        if line.startswith(ignored):
            continue

        # Check line length
        line = "{:<80}\n".format(line.strip())
        
        yield line
        
        if line.startswith(records):
            prev_line = line
            break

    # Now go through all the remaining lines
    serial_equiv = {}  # store for conect statements
    coord_section = False
    serial_offset = 0  # To offset after adding TER records
    for line in fhandle:
        if line.startswith(ignored):
            continue

        elif line.startswith(records):
            coord_section = True
            is_gap = (int(line[22:26]) - int(prev_line[22:26])) > 1
            if line[21] != prev_line[21] or is_gap:
                serial = int(prev_line[6:11]) + 1
                rname = prev_line[17:20]
                chain = prev_line[21]
                resid = prev_line[22:26]
                icode = prev_line[26]
                
                serial_offset += 1
                
                ter_line = fmt_TER.format(serial, rname, chain, resid, icode)
                yield ter_line

            ori_serial = int(prev_line[6:11])
            serial = ori_serial + serial_offset

            serial_equiv[ori_serial] = serial

            line = line[:6] + str(serial).rjust(5) + line[11:]
            prev_line = line
            serial_offset = 1  # reset otherwise we have +2 ex.

        elif line.startswith('ANISOU'):
            # Fix serial based on previous atom
            serial = prev_line[6:11]
            line = line[:6] + str(serial).rjust(5) + line[11:]

        elif line.startswith('CONECT'):
            # 6:11, 11:16, 16:21, 21:26, 26:31
            serials = [line[cr] for cr in char_ranges]
            # If not found, return default
            new_serials = [serial_equiv.get(s, s) for s in serials]
            conect_line = fmt_CONECT.format(*new_serials)
            yield conect_line
            continue

        else:
            if coord_section:
                coord_section = False
                # Add last TER statement
                serial = int(prev_line[6:11]) + 1
                rname = prev_line[17:20]
                chain = prev_line[21]
                resid = prev_line[22:26]
                icode = prev_line[26]

                ter_line = fmt_TER.format(serial, rname, chain, resid, icode)
                yield ter_line
        
        # Check line length
        line = "{:<80}\n".format(line.strip())
        
        yield line
    
    # Add END statement
    yield "{:<80}\n".format("END")


def main():
    # Check Input
    pdbfh = check_input(sys.argv[1:])
{'fraction_non_alphanumeric': 0.10295320308950477, 'fraction_numerical': 0.03144025442980464, 'mean_word_length': 2.0386526780784098, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2077271', 'n_tokens_mistral': 3626, 'n_tokens_neox': 3290, 'n_words': 1060}