This script runs black on a text file with doctest and blackens them....

import argparse import sys

import black from blib2to3.pgen2.tokenize import TokenError

TEST_DATA = """ Normal

name = 'matt' for num in [1,2,3]: ... print(num) ... print(2, num)

Rest

"""

LONG_LINES = '''

NORMAL

def draw_cmap(name, size=10, aspect=.25): ... fig=plt.figure() ... fig.set_size_inches(4,4) ... #ax = plt.subplot(111) ... ax = plt.axes([0,0,1,1], frameon=False) ... mapname = name ... set_cmap(mapname) ... colors = getattr(cm, mapname) ... # Then we disable our xaxis and yaxis completely. If we just say plt.axis('off'), ... # they are still used in the computation of the image padding.

END '''

def line_class(fin): r""" >>> for data in line_class(TEST_DATA.split('\n')): ... print(data) (3, '') (3, 'Normal') (3, '') (1, ">>> name = 'matt'") (1, '>>> for num in [1,2,3]:') (2, '... print(num)') (2, '... print(2, num)') (3, '') (3, 'Rest') (3, '') (3, '') (3, '') """ for line in fin: sline = line.strip() if sline.startswith('>>>'): klass = 1 elif sline.startswith('...'): klass = 2 else: klass = 3 yield klass, line

def code_chunks(fin): r""" returns tuples of (klass, [lines]) klass: 1 is console code 2 non-console code

>>> lines = TEST_DATA.split('\n')
>>> print(lines)
['', 'Normal', '', ">>> name = 'matt'", '>>> for num in [1,2,3]:', '...     print(num)', '...     print(2, num)', '', 'Rest', '', '', '']
>>> print((list(code_chunks(lines))))
[(2, ['']), (2, ['Normal']), (2, ['']), (1, [">>> name = 'matt'"]), (1, ['>>> for num in [1,2,3]:', '...     print(num)', '...     print(2, num)']), (2, ['']), (2, ['Rest']), (2, ['']), (2, ['']), (2, [''])]

"""
in_code = False
chunk = []
for klass, line in line_class(fin):
    if in_code:
        if klass == 1:
            if chunk:
                yield 1, chunk  # previously saw >>>, yield and add >>> in own chunk
            chunk = [line]
        elif klass == 2:      # add ... to chunk
            chunk.append(line)
        elif klass == 3:
            if chunk:
                yield 1, chunk
            yield 2, [line]
            chunk = []
    else:
        if klass == 1:
            chunk = [line]
            in_code = True
        elif klass == 2:
            print("ERROR!")
        else:
            if chunk:
                yield 1, chunk
                chunk = []
            in_code = False
            yield 2, [line]
if chunk:
    yield 1, chunk
      

def test_process(): r""" >>> import io >>> out = io.StringIO() >>> data = [f'{line}\n' for line in LONG_LINES.split('\n')] >>> process(data, out, chars=30) >>> print(out.getvalue()) NORMAL >>> def draw_cmap( ... name, ... size=10, ... aspect=0.25, ... ): ... fig = plt.figure() ... fig.set_size_inches( ... 4, 4 ... ) ... # ax = plt.subplot(111) ... ax = plt.axes( ... [0, 0, 1, 1], ... frameon=False, ... ) ... mapname = name ... set_cmap(mapname) ... colors = getattr( ... cm, mapname ... ) ... # Then we disable our xaxis and yaxis completely. If we just say plt.axis('off'), ... # they are still used in the computation of the image padding. END """ pass

def process(fin, fout, chars=51): lines = [] add_newlines = True joiner = '\n' orig_line_num = 0 new_line_num = 0 for i, (klass, chunk) in enumerate(code_chunks(fin)): if i == 0: first = chunk[0] add_newlines = not first.endswith('\n') if not add_newlines: joiner = '' if klass == 2: lines.append(chunk[0]) else: first = chunk[0] whitespace_len = len(first) - len(first.lstrip()) new_chunk = [line[whitespace_len + 4:] for line in chunk] mode = black.FileMode(line_length=chars-4) #target_versions=set(), try: new_content = black.format_str(joiner.join(new_chunk), mode=mode) #new_content = black.format_file_contents(joiner.join(new_chunk), # fast=True, mode=mode) except black.InvalidInput as ex: print("LINENUM", orig_line_num) print("ERRR!", joiner.join(new_chunk)) raise except TokenError: print("2LINENUM", orig_line_num) print("2ERRR!", joiner.join(new_chunk)) raise new_with_prompts = new_content.split('\n') new_with_prompts[0] = f'{" "*whitespace_len}>>> {new_with_prompts[0]}' for i in range(1, len(new_with_prompts)): new_with_prompts[i] = f'{" "*whitespace_len}... {new_with_prompts[i]}' if new_with_prompts[-1].strip() == '...': new_with_prompts = new_with_prompts[:-1] lines.extend([f'{line}\n' for line in new_with_prompts]) new_line_num = len(lines) orig_line_num += len(chunk) fout.write(joiner.join(lines))

def main(args): ap = argparse.ArgumentParser(description='Look for python console snippets and apply black to them') ap.add_argument('-s', '--src', help='src file') ap.add_argument('-d', '--dst', help='dst file (default stdout)', default=sys.stdout) ap.add_argument('-l', '--length', help='black split column (from >>> position, so if you start >>> at col 15 and set this to 30 you only have til column 45) default 51', type=int, default=51) ap.add_argument('-t', '--test', help='run doctest', action='store_true') opt = ap.parse_args(args) if opt.src:

    with open(opt.src) as fin:
        if opt.dst != sys.stdout:
            fout = open(opt.dst, 'w')
        else:
            fout = sys.stdout
        process(fin, fout, chars=opt.length)
if opt.test:
    import doctest
    doctest.testmod()

if name == 'main': main(sys.argv[1:])