#!/bin/bash
  
# usage: nbgrep 'pattern'

SEARCHPATH='/Users/andrew.blevins/projects'

# 'jq' technique lifted with gratitude
# from https://gist.github.com/mlgill/5c55253a3bc84a96addf
# from https://gist.github.com/jbarratt/fa1d3473048e5f856aeb

# Break on newlines instead of any whitespace
# IPython Notebook files often have spaces in it
SAVEIFS=$IFS
IFS=$(echo -en "\n\b")

# mdfind uses OSX's spotlight search, so it's almost instant
# generate a list of all the ipynb files in any of the directories
FILES=$(mdfind -onlyin $SEARCHPATH -name '.ipynb')

# On the command line we get the argument to search for
PATTERN=$1

for f in $FILES
do
    # Check Notebook JSON format first
       # Jupyter notebook JSON format
        OUTPUT=$(jq '.cells[]? | select(.cell_type=="code") | .source[]?//.source' $f \
            | sed 's/^"//g;s/"$//g;s/\\n$//g;s/\\"/"/g;s/\\\\/\\/g;s/\\n/\n/g' \
            | pygmentize -l python 2>/dev/null \
            | grep $PATTERN 2>/dev/null)

    if [ $? -eq 0 ]; then
        echo -e "$f:\n\n$OUTPUT\n\n"
    fi
done

IFS=$SAVEIFS