Python insert line in the middle of the file

Question

I will be inserting a new text line in a setup.py file, this new line will contain text, and is part of the REQUIRED_PACKAGES. I parsed the file and look for REQUIRED_PACKAGES pattern (REQUIRED_PACKAGES = [ ... ]). Looking for alternatives during file insertion.

A sample setup.py file looks like this:

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import fnmatch
import os
import re
import sys

from setuptools import Command
from setuptools import find_packages
from setuptools import setup
from setuptools.command.install import install as InstallCommandBase
from setuptools.dist import Distribution

DOCLINES = __doc__.split('\n')

_VERSION = '1.0.0'

REQUIRED_PACKAGES = [
    'absl-py >= 0.7.0',
    'astunparse == 1.6.3',
    'backports.weakref >= 1.0rc1;python_version<"3.4"',
    'enum34 >= 1.1.6;python_version<"3.4"',
    'gast == 0.3.3',
    'scipy == 1.2.2;python_version<"3"',
]

if sys.byteorder == 'little':
  # grpcio does not build correctly on big-endian machines due to lack of
  # BoringSSL support.  
  REQUIRED_PACKAGES.append('io >= 1.0.0')

...

After conversion it will look like this (adding a new line):

...
REQUIRED_PACKAGES = [
    'absl-py >= 0.7.0',
    'astunparse == 1.6.3',
    'backports.weakref >= 1.0rc1;python_version<"3.4"',
    'enum34 >= 1.1.6;python_version<"3.4"',
    'gast == 0.3.3',
    'scipy == 1.2.2;python_version<"3"',
    'test==1.0.0',
]
...

Python code

import re


SEARCH_DICT = {
    'required_packages': re.compile(
        r'REQUIRED_PACKAGES = (?P<required_packages>.*)\n')
}

TEST_LIBRARY = '\t\t\'test==1.0.0\'\n'


def _parse_line(line):
    """
    Do a regex search against all defined regexes and
    return the key and match result of the first matching regex

    """

    for key, rx in SEARCH_DICT.items():
        match = rx.search(line)
        if match:
            return key, match
    # if there are no matches
    return None, None


def parse_file(filepath):
    """
    Parse text at given filepath

    Parameters
    ----------
    filepath : str
        Filepath for file_object to be parsed

    Returns
    -------
    data : file contents

    while True:
        last_pos = fp.tell()
        line = fp.readline()
        new_line = line.upper()
        fp.seek(last_pos)
        fp.write(new_line)
        print "Read %s, Wrote %s" % (line, new_line)


    """

    data = []  # create an empty list to collect the data
    line_index = -1 
    # open the file and read through it line by line

    with open(filepath, 'r+') as file_object:
        line = file_object.readline()
        line_index+=1
        while line:
            # at each line check for a match with a regex
            key, match = _parse_line(line)
            if key == 'required_packages':
                required_packages_start = match.group('required_packages')
                if required_packages_start == '[':                    
                    print('Found REQUIRED_PACKAGES')
                    while line.strip():                        
                        library = line.rstrip()
                        if library == ']': # End of required packages
                            return line_index
                        line = file_object.readline()
                        line_index+=1                     

            line = file_object.readline()  
            line_index+=1          
        file_object.readline()
        line_index+=1

    return line_index


line_index = parse_file('setup.py')
lines = None
if line_index != -1:
    with open('setup.py', 'r+') as fd:
        contents = fd.readlines()
        contents.insert(line_index, TEST_LIBRARY)  # new_string should end in a newline
        fd.seek(0)  # readlines consumes the iterator, so we need to start over
        fd.writelines(contents)  # No need to truncate as we are increasing filesize

Answer 1

Your code goes through the file twice - once to find the index where your new line is to be inserted, and again to actually perform the change.

My approach to this problem would be to create a class that contains three variables:

1. The lines in setup.py before "REQUIRED_PACKAGES = ["
2. The contents of the REQUIRED_PACKAGES list
3. The lines in setup.py after the REQUIRED_PACKAGES list

Your problem then splits nicely into three distinct sub-problems:

1. Parsing setup.py and returning an instance of the class above
2. Adding an element to the required packages list in your custom class (trivial)
3. Writing a new setup.py file

Looking at your code specifically, your docstring for parse_file does not match what the code is doing. It says that data is returned, but the code actually returns line_index, and the data member you initialise is unused.

Your code comments in general could be improved. You comment some obvious things such as"open the file and read through it line by line" and "if there are no matches". These operations are immediately apparent from looking at the code, and do not require comments.

Answer 2

Things I would do differently:

• make the function checking for the line where to insert accept any iterator which yields lines. This hoists the IO to a level above, and makes the function easier to test. I would also throw an exception if there is no place to insert the line, instead of returning -1 as a sentinel value
• make the function inserting the line a generator that yields lines
• use a temporary file to write the new content, and when finished, copy this to the original file

Finding the insertion place

With some rearrangements, replacing the while for a for-loop and inversions of checks, I arrive at something like this

def find_insert_position(lines: typing.Iterator[str]) -> int:
    line_index = -1
    for line in lines:
        line_index += 1

        key, match = _parse_line(line)
        if key != "required_packages":
            continue
        required_packages_start = match.group("required_packages")
        if required_packages_start != "[":
            continue
        logging.debug("Found REQUIRED_PACKAGES")
        while True:
            library = line.rstrip()
            if library == "]":  # End of required packages
                return line_index
            try:
                line = next(lines)
                line_index += 1
            except StopIteration:
                raise ValueError("No required packages found")
    raise ValueError("No required packages found")

I also started using the logging module instead of a print, and added type annotations.

I also adapted the regex a bit to

SEARCH_DICT = {
    "required_packages": re.compile(
        r"^REQUIRED_PACKAGES = (?P<required_packages>.*)$"
    )
}

This works in both your and this method.

In your original method, I found no reason for the while line.strip(). On the contrary. If there is an empty line somewhere in the list of required packages, this does not work. I changed this to a while True.

This code can be tested with a number of testcases, listed at the end of the answer

test_cases = {
    test_case_empty: None,
    test_case_no_req: None,
    test_case_empty_req: 2,
    test_case_one_line_req: None,
    test_case_req: 8,
    code: 27,
    test_case_req_empty_line: 9,
}

for test_case, expected_result in test_cases.items():
    try:
        result = find_insert_position(test_case.split("\n"))
        assert result == expected_result
    except ValueError as e:
        assert expected_result is None

inserting the line

This can be very easily done with a generator:

import itertools
def insert_line(lines: typing.Iterable[str], insert_position: int, line_to_insert: str):
    lines = iter(lines)
    for i, line in enumerate(itertools.islice(lines, insert_position)):
        yield line
    if insert_position > i + 1:
        raise ValueError("insert_position is beyond `lines`")
    yield line_to_insert
    yield from lines

And tested as follows:

list(insert_line(["a"] * 10, 10, "test"))

['a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'a', 'test']

list(insert_line(["a"] * 10, 11, "test"))

Raises ValueError

So this is also a simple, easily testable unit.

tempfile

Instead of reading and writing to the same file, I would either create a temporary file (with tempfile and then replace the original with it, or concatenate the results into a string and then write that to the file

def write_result(lines, output_file: Path):
    with tempfile.TemporaryDirectory() as tempdir:
        my_tempfile : Path = Path(tempdir) / "output.txt"
        with my_tempfile.open("w") as filehandle:
            for line in lines:
                filehandle.write(line)
        my_tempfile.replace(output_file)

or

result = "\n".join(lines)
ast.parse(result)
outut_file.write_text(result)

The main difference is when the source is large, this will use more memory

verify the result

You can even use the ast or py_compile module to check whether you generated valid python by inserting a py_compile.compile(my_tempfile) before replacing the original with the generate file

putting it together:

if __name__ == "__main__":
    source_file = Path("test.py")
    TEST_LIBRARY = "    'test,'\n"
    with source_file.open("r") as filehandle:
        try:
            insert_location = find_insert_position(filehandle)
        except ValueError:
            insert_location = None
    print(f"insert_location: {insert_location}")
    if insert_location is not None:
        with source_file.open("r") as filehandle:
            new_lines = list(insert_line(
                lines=filehandle,
                insert_position=insert_location,
                line_to_insert=TEST_LIBRARY,
            ))
        write_result(new_lines, source_file.with_suffix(".new.py"))

I've added the .with_suffix(".new.py") so during testing, you can check the results side by side.

I've also replaced the ' in TEST_LIBRARY with " so you don't have to escape the '. I also see you use tabs to indent, better use 4 spaces.

Also know that this approach might not work when the last line in the REQUIRED_PACKAGES does not end with a ,. Python concatenates strings on different lines, so this will not be picked up by the py_compile either, so this will fail silently.

---

test cases

test_case_empty = """


"""
test_case_no_req = """
no required packages
"""

test_case_empty_req = """
REQUIRED_PACKAGES = [
]
"""

test_case_one_line_req = """
REQUIRED_PACKAGES = []
"""

test_case_req = """
REQUIRED_PACKAGES = [
    'absl-py >= 0.7.0',
    'astunparse == 1.6.3',
    'backports.weakref >= 1.0rc1;python_version<"3.4"',
    'enum34 >= 1.1.6;python_version<"3.4"',
    'gast == 0.3.3',
    'scipy == 1.2.2;python_version<"3"',
]
"""
test_case_req_empty_line = """
REQUIRED_PACKAGES = [
    'absl-py >= 0.7.0',
    'astunparse == 1.6.3',
    'backports.weakref >= 1.0rc1;python_version<"3.4"',
    'enum34 >= 1.1.6;python_version<"3.4"',

    'gast == 0.3.3',
    'scipy == 1.2.2;python_version<"3"',
]
"""