py-to-ipynb.py

#!/usr/bin/env python3
# Copyright kurinoku 2023
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy of 
# this software and associated documentation files (the “Software”), to deal 
# in the Software without restriction, including without limitation the rights to 
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 
# of the Software, and to permit persons to whom the Software is furnished to do 
# so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all 
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 
# SOFTWARE.

"""
py-to-ipynb.py ifile ofile
  ifile: input .py file
  ofile: output .ipynb file

the syntax for the ifile is the following
=== main.py ===
#> markdown
# # My hello world
# You have to prepend each markdown line
# with a # so it is a valid python file
#> python
print('Hello World')
#> python
#>! matplotlib inline
#>% pip install matplotlib
== END ===

This code is not really tested, so use it at your own risk.
I also don't really use notebooks that much, so it could
be lacking features I don't know, also every time you regenerate,
each cell has a new id, I guess it's used for "checkpoints".
It probably breaks that stuff.
"""

import json
import traceback
import sys
import argparse

make_python_cell = lambda id, source: {
   "cell_type": "code",
   "execution_count": 0,
   "id": id,
   "metadata": {},
   "outputs": [],
   "source": source
  }

make_md_cell = lambda id, source: {
   "cell_type": "markdown",
   "id": id,
   "metadata": {},
   "source": source
  }

make_notebook = lambda cells: {
 "cells": cells,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}

ID_START = 0x11111111
as_id = lambda x: '{:x}'.format(x)

CELL_TYPE_PY = 'python'
CELL_TYPE_MD = 'markdown'

class Cell:
    def __init__(self):
        self.sources = []
    def add(self, source):
        source = self._process_line(source)
        if source is not None:
            self.sources.append(source)

    def _process_line(self, line):
        raise NotImplementedError
    def make_nb(self, id):
        raise NotImplementedError

class ProcessLineError(Exception):
    pass

class PyCell(Cell):
    def _process_line(self, line):
        og_line = line
        line = line.strip()
        if line.startswith('#>!') or line.startswith('#>%'):
            c = line[2]
            rest = line[3:].strip()
            if not rest:
                raise ProcessLineError(f'expected a command after "{c}"')
            return c + ' ' + rest
        return og_line.rstrip()

    def make_nb(self, id):
        return make_python_cell(id, [line + '\n' for line in self.sources])

class MdCell(Cell):
    def _process_line(self, line):
        line = line.strip()
        if not line:
            return None
        if line.startswith('# ') or line == '#':
            return line[2:]
        raise ProcessLineError('Md does not understand line')
    
    def make_nb(self, id):
        return make_md_cell(id, [line + '\n' for line in self.sources])

LANGS = [('python', PyCell), ('markdown', MdCell)]

def parse(s):
    lines = s.split('\n')
    cells = []
    current_cell = None
    line_count = 0
    for line in lines:
        if line.strip() == '#>':
            print(f'at line {line_count}: syntax error', file=sys.stderr)
            sys.exit(-1)
        elif line.startswith('#> '):
            _, language = line.split(' ', maxsplit=1)
            language = language.strip().lower()
            for (lang, klass) in LANGS:
                if language in lang:
                    if current_cell is not None:
                        cells.append(current_cell)
                    current_cell = klass()
                    break
            else:
                print(f'at line {line_count}: language "{language}" not recognized', file=sys.stderr)
                sys.exit(-1)
        else:
            if not line.strip():
                continue
            try:
                if current_cell is None:
                    raise ProcessLineError('tried to read line when not inside a cell')
                current_cell.add(line)
            except ProcessLineError as exn:
                print(f'at line {line_count}: exception found', file=sys.stderr)
                traceback.print_exc(exn)
                sys.exit(-1)


        line_count += 1
    if current_cell is not None:
        cells.append(current_cell)
    return cells

def write_notebook(ofile, cells):
    id = ID_START
    final_cells = []
    for c in cells:
        t = c.make_nb(as_id(id))
        final_cells.append(t)
        id += 0x1

    o = make_notebook(final_cells)
    json.dump(o, ofile)

def main():
    a = argparse.ArgumentParser()
    a.add_argument('ifile')
    a.add_argument('ofile')
    args = a.parse_args()
    with open(args.ifile, 'r', encoding='utf-8') as ifile:
        processed = parse(ifile.read())

    with open(args.ofile, 'w', encoding='utf-8') as ofile:
        write_notebook(ofile, processed)

if __name__ == '__main__':
    main()