# AUTOGENERATED! DO NOT EDIT! File to edit: ../notebooks/nb_definition.ipynb.

# %% auto 0
__all__ = ['DEFAULT_NB_METADATA', 'Notebook', 'CleanOperations']

# %% ../notebooks/nb_definition.ipynb 3
from nbformat import NotebookNode
from nbformat import read as read_nb, write as write_nb, validate as validate_nb, from_dict
import jq

# %% ../notebooks/nb_definition.ipynb 5
from typing import Iterable
from enum import Enum

DEFAULT_NB_METADATA = '{"kernelspec": {"language": "python"}}'

# TODO: about filter, how about:
# nb.cells = jq.compile(filter).input(nb.cells).all()?
class Notebook:
    def __init__(self, path):
        self.path = path
        self.nb = read_nb(path, as_version=4)

    def apply_jq_pattern(self, pattern):
        return from_dict(jq.compile(pattern).input(self.nb).first())

    def clean_execution_counts(self):
        self.nb = self.apply_jq_pattern('.cells[].execution_count = null')
        
    def clean_outputs(self):
        self.nb = self.apply_jq_pattern('.cells[].outputs = []')
    
    def clean_all_metadata(self):
        self.nb = self.apply_jq_pattern('.cells[].metadata = {}')
        self.nb = self.apply_jq_pattern(f'.metadata = {DEFAULT_NB_METADATA}')
    
    def clean_all(self):
        self.clean_execution_counts()
        self.clean_outputs()
        self.clean_all_metadata()
        
    # TODO: circular dependency - notebooks.clean needs clean operations and vice versa
    def clean(self, clean_ops: Iterable[Enum]):        
        for op in clean_ops:
            getattr(self, op.value)()
        return self
        
    def save(self, path=None):
        path = path if path else self.path
        write_nb(self.nb, path)
        

# %% ../notebooks/nb_definition.ipynb 6
class CleanOperations(Enum):
    METADATA = Notebook.clean_all_metadata.__name__
    OUTPUTS = Notebook.clean_outputs.__name__
    EXECUTION_COUNT = Notebook.clean_execution_counts.__name__
    ALL = Notebook.clean_all.__name__ 

