Gnarly code extractor for Django Crash Course

extractor.py

LEGALESE = """Using This Code Example
=========================
The code examples provided are provided by Daniel and Audrey Feldroy of
Feldroy to help you reference Django Crash Course. Code samples follow
PEP-0008, with exceptions made for the purposes of improving book
formatting. Example code is provided "as is".

Permissions
============
In general, you may use the code we've provided with this book in your
programs . You do not need to contact us for permission unless you're
reproducing a significant portion of the code and using it in educational
distributions. Examples:
* Writing an education program or book that uses several chunks of code from
    this course requires permission.
* Selling or distributing a digital package from material taken from this
    book does require permission.
* Answering a question by citing this book and quoting example code does not
    require permission.
* Incorporating a significant amount of example code from this book into your
    product's documentation does require permission.
Attributions usually include the title, author, publisher and an ISBN. For
example, "Django Crash Course, by Daniel and
Audrey Feldroy. Copyright 2020 Feldroy."

If you feel your use of code examples falls outside fair use of the permission
given here, please contact us at [email protected].
"""

import os
import shutil
from glob import glob

try:
    shutil.rmtree("code")
    print("Removed old code directory")
except FileNotFoundError:
    pass
os.mkdir("code")
print("Created new code directory")

files = glob("[0-9]*.Rmd")
files.sort()


def code_check(line):
    if not line.startswith("```"):
        return False
    return "block2" not in line and "knitr-logo" not in line


def commentator(code_type):
    if code_type == "python":
        return f'"""\n{LEGALESE}"""'
    if code_type.strip() in "bash,plaintext,text,ini,conf":
        legalese = [f"# {x}" for x in LEGALESE.splitlines()]
        return "\n".join(legalese)
    if code_type == "djangotemplate":
        return f'<!--\n{LEGALESE}-->'
    return LEGALESE


TYPES = {
    "python": "py",
    "djangotemplate": "html",
    "plaintext": "txt",
    "text": "txt",
    "ini": "ini",
    "conf": "conf",
    "bash": "bash",
}


for filename in files:
    prefix = filename.split("-")[0]

    in_code = False
    code = []
    example_num = 0

    with open(filename) as f:

        for line in f.readlines():
            if code_check(line) and in_code == False:
                code_type = line[3:].strip()
                if not len(code_type.strip()):
                    continue
                extension = TYPES[code_type]
                example_num += 1
                in_code = True
                continue

            if line.startswith("```") and in_code == True:
                new_filename = (
                    f"code/chapter_{prefix}_example_{str(example_num).zfill(2)}.{extension}"
                )
                # print(new_filename)
                with open(new_filename, "w") as f:
                    f.write(commentator(code_type))
                    f.write("\n\n")
                    f.write("".join(code))
                in_code = False
                code = []
                continue

            if in_code:
                code.append(line)