<commit_before>#!/usr/bin/env python
"""
master_parse.py: convert master Databricks notebooks to other formats

Converts notebooks exported in "source" format from Databricks
into separate notebooks by language and for instructor and student use.
It can also generate IPython notebooks.

Works with Python 2 and Python 3. When using IPython, IPython version 3 is
required.

This module can be used as a command or a library. To see the valid command
line arguments, run with no arguments. If used as a library, the main
entry point is process_notebooks().
"""

from __future__ import print_function

import glob
import os
import os.path
import re
import codecs
from enum import Enum

VERSION = "1.2.0"

CommandLabel = Enum('CommandLabel',
                    'IPYTHON_ONLY PYTHON_ONLY SCALA_ONLY ' +
                    'R_ONLY SQL_ONLY ANSWER TODO TEST PRIVATE_TEST ' +
                    'DATABRICKS_ONLY INLINE ALL_NOTEBOOKS INSTRUCTOR_NOTES')
CommandCode = Enum('CommandCode',
                   'SCALA PYTHON R SQL MARKDOWN FILESYSTEM SHELL RUN FS SH')
NotebookKind = Enum('NotebookKind', 'DATABRICKS IPYTHON')
NotebookUser = Enum('NotebookUser', 'INSTRUCTOR STUDENT')
NewlineAfterCode = {CommandCode.SCALA, CommandCode.R, CommandCode.PYTHON}

from collections import namedtuple
Command = namedtuple('Command', ['part', 'code', 'labels', 'content'])


class NotebookGenerator(object):
    code_to_output_dir = {CommandCode.SQL:    'sql',
                          CommandCode.SCALA:  'scala',
                          CommandCode.PYTHON: 'python',
                          CommandCode.R:      'r'}

    param_to_label = {
        CommandCode.PYTHON: {CommandLabel.PYTHON_ONLY},
        CommandCode.SCALA: {CommandLabel.SCALA_ONLY},
        CommandCode.SQL: {CommandLabel.SQL_ONLY},
        CommandCode.R: {CommandLabel.R_ONLY},
        CommandCode.FS: {CommandLabel.ALL_NOTEBOOKS},
        CommandCode.SH: {CommandLabel.ALL_NOTEBOOKS},
        NotebookKind.DATABRICKS: {CommandLabel.DATABRICKS_ONLY},
        NotebookKind.IPYTHON: {CommandLabel.IPYTHON_ONLY},
        NotebookUser.INSTRUCTOR: {CommandLabel.ANSWER,
                                  CommandLabel.INSTRUCTOR_NOTES,
                                  CommandLabel.PRIVATE_TEST},
        NotebookUser.STUDENT: {CommandLabel.TODO},
    }

    user_to_extension = {
        NotebookUser.INSTRUCTOR: '_answers',
        NotebookUser.STUDENT: '_student'
    }

    def __init__(self, notebook_kind, notebook_user, notebook_code):
        base_keep = set(CommandLabel.__members__.values())
        self.keep_labels = self._get_labels(notebook_kind,
                                            notebook_user,
                                            notebook_code)
        # discard labels not explicitly kept
        self.discard_labels = base_keep - self.keep_labels
        self.remove = [_dbc_only, _scala_only, _python_only, _new_part, _inline,
                       _all_notebooks, _instructor_notes]
        self.replace = [(_ipythonReplaceRemoveLine, ''),
                        _rename_public_test,
                        _rename_import_public_test]
        self.notebook_kind = notebook_kind
        self.notebook_user = notebook_user
        self.notebook_code = notebook_code
        self.file_ext = self._get_extension()
        self.base_comment = _code_to_comment[self.notebook_code]

    def _get_labels(self, *params):
        labels = {CommandLabel.TEST, CommandLabel.INLINE}
        for param in params:
            labels.update(NotebookGenerator.param_to_label[param])
        return labels

    def _get_extension(self):
        return (NotebookGenerator.user_to_extension[self.notebook_user] +
                Parser.code_to_extension[self.notebook_code])

    def generate(self, header, commands, input_name, parts=True,
                 creative_commons=False):
        is_IPython = self.notebook_kind == NotebookKind.IPYTHON

        command_cell = _command_cell.format(self.base_comment)

        max_part = max([part for (part, _, _, _) in commands]) if parts else 0
        if max_part == 0 or is_IPython:
            parts = False

        # generate full notebook when generating parts
        if parts:
            self.generate(header, commands, input_name, False,
                          creative_commons=creative_commons)

        base_file = os.path.splitext(os.path.basename(input_name))[0]

        for i in range(max_part + 1): # parts are zero indexed

            part_base = '_part{0}'
            part_string = part_base.format(i + 1) if parts else ''

            out_dir = os.path.join(
                _output_dir, base_file,
                NotebookGenerator.code_to_output_dir[self.notebook_code]
            )

            if not os.path.exists(out_dir):
                os.makedirs(out_dir)

            file_out = os.path.join(out_dir,
                                    base_file + part_string + self.file_ext)
            if is_IPython:
                file_out = file_out.replace('.py', '.ipynb')

            with codecs.open(file_out, 'w', _file_encoding_out,
                             errors=_file_encoding_errors) as output:

                # don't write the first # command --- for databricks
                is_first = not is_IPython

                if is_IPython:
                    header_adj = _header
                else:
                    # use proper comment char
                    header_adj = re.sub(_comment, self.base_comment, header)
                    if creative_commons:
                        header_adj += '\n{0} MAGIC %md\n{0} MAGIC {1}'.format(
                            self.base_comment, _cc_license
                        )
                        is_first = False

                output.write(header_adj + '\n')

                added_run = False
                for (part, code, labels, content) in commands:
                    if parts:
                        if part > i:  # don't show later parts
                            break
                        elif part == i - 1 and not added_run:
                            # add %run command before new part
                            student = NotebookGenerator.user_to_extension[
                                NotebookUser.STUDENT
                            ]
                            instr = NotebookGenerator.user_to_extension[
                                NotebookUser.INSTRUCTOR
                            ]
                            solution_file = file_out.replace(student, instr)

                            solution_file = os.path.splitext(solution_file)[0]
                            solution_file = solution_file.replace(
                                part_string, part_base.format(i)
                            )
                            db_location = ('/Users/<EMAIL>/Labs/' +
                                           solution_file.replace('\\', '/'))
                            db_location = db_location.replace('build_mp/', '')
                            runCommand = '{0} MAGIC %run {1}'.format(
                                self.base_comment, db_location
                            )
                            is_first = self._write_command(
                                output, command_cell, [runCommand, ''], is_first<commit_msg>%run is now relative, as supported by DB.<commit_after>#!/usr/bin/env python
"""
master_parse.py: convert master Databricks notebooks to other formats

Converts notebooks exported in "source" format from Databricks
into separate notebooks by language and for instructor and student use.
It can also generate IPython notebooks.

Works with Python 2 and Python 3. When using IPython, IPython version 3 is
required.

This module can be used as a command or a library. To see the valid command
line arguments, run with no arguments. If used as a library, the main
entry point is process_notebooks().
"""

from __future__ import print_function

import glob
import os
import os.path
import re
import codecs
from enum import Enum

VERSION = "1.2.0"

CommandLabel = Enum('CommandLabel',
                    'IPYTHON_ONLY PYTHON_ONLY SCALA_ONLY ' +
                    'R_ONLY SQL_ONLY ANSWER TODO TEST PRIVATE_TEST ' +
                    'DATABRICKS_ONLY INLINE ALL_NOTEBOOKS INSTRUCTOR_NOTES')
CommandCode = Enum('CommandCode',
                   'SCALA PYTHON R SQL MARKDOWN FILESYSTEM SHELL RUN FS SH')
NotebookKind = Enum('NotebookKind', 'DATABRICKS IPYTHON')
NotebookUser = Enum('NotebookUser', 'INSTRUCTOR STUDENT')
NewlineAfterCode = {CommandCode.SCALA, CommandCode.R, CommandCode.PYTHON}

from collections import namedtuple
Command = namedtuple('Command', ['part', 'code', 'labels', 'content'])


class NotebookGenerator(object):
    code_to_output_dir = {CommandCode.SQL:    'sql',
                          CommandCode.SCALA:  'scala',
                          CommandCode.PYTHON: 'python',
                          CommandCode.R:      'r'}

    param_to_label = {
        CommandCode.PYTHON: {CommandLabel.PYTHON_ONLY},
        CommandCode.SCALA: {CommandLabel.SCALA_ONLY},
        CommandCode.SQL: {CommandLabel.SQL_ONLY},
        CommandCode.R: {CommandLabel.R_ONLY},
        CommandCode.FS: {CommandLabel.ALL_NOTEBOOKS},
        CommandCode.SH: {CommandLabel.ALL_NOTEBOOKS},
        NotebookKind.DATABRICKS: {CommandLabel.DATABRICKS_ONLY},
        NotebookKind.IPYTHON: {CommandLabel.IPYTHON_ONLY},
        NotebookUser.INSTRUCTOR: {CommandLabel.ANSWER,
                                  CommandLabel.INSTRUCTOR_NOTES,
                                  CommandLabel.PRIVATE_TEST},
        NotebookUser.STUDENT: {CommandLabel.TODO},
    }

    user_to_extension = {
        NotebookUser.INSTRUCTOR: '_answers',
        NotebookUser.STUDENT: '_student'
    }

    def __init__(self, notebook_kind, notebook_user, notebook_code):
        base_keep = set(CommandLabel.__members__.values())
        self.keep_labels = self._get_labels(notebook_kind,
                                            notebook_user,
                                            notebook_code)
        # discard labels not explicitly kept
        self.discard_labels = base_keep - self.keep_labels
        self.remove = [_dbc_only, _scala_only, _python_only, _new_part, _inline,
                       _all_notebooks, _instructor_notes]
        self.replace = [(_ipythonReplaceRemoveLine, ''),
                        _rename_public_test,
                        _rename_import_public_test]
        self.notebook_kind = notebook_kind
        self.notebook_user = notebook_user
        self.notebook_code = notebook_code
        self.file_ext = self._get_extension()
        self.base_comment = _code_to_comment[self.notebook_code]

    def _get_labels(self, *params):
        labels = {CommandLabel.TEST, CommandLabel.INLINE}
        for param in params:
            labels.update(NotebookGenerator.param_to_label[param])
        return labels

    def _get_extension(self):
        return (NotebookGenerator.user_to_extension[self.notebook_user] +
                Parser.code_to_extension[self.notebook_code])

    def generate(self, header, commands, input_name, parts=True,
                 creative_commons=False):
        is_IPython = self.notebook_kind == NotebookKind.IPYTHON

        command_cell = _command_cell.format(self.base_comment)

        max_part = max([part for (part, _, _, _) in commands]) if parts else 0
        if max_part == 0 or is_IPython:
            parts = False

        # generate full notebook when generating parts
        if parts:
            self.generate(header, commands, input_name, False,
                          creative_commons=creative_commons)

        base_file = os.path.splitext(os.path.basename(input_name))[0]

        for i in range(max_part + 1): # parts are zero indexed

            part_base = '_part{0}'
            part_string = part_base.format(i + 1) if parts else ''

            out_dir = os.path.join(
                _output_dir, base_file,
                NotebookGenerator.code_to_output_dir[self.notebook_code]
            )

            if not os.path.exists(out_dir):
                os.makedirs(out_dir)

            file_out = os.path.join(out_dir,
                                    base_file + part_string + self.file_ext)
            if is_IPython:
                file_out = file_out.replace('.py', '.ipynb')

            with codecs.open(file_out, 'w', _file_encoding_out,
                             errors=_file_encoding_errors) as output:

                # don't write the first # command --- for databricks
                is_first = not is_IPython

                if is_IPython:
                    header_adj = _header
                else:
                    # use proper comment char
                    header_adj = re.sub(_comment, self.base_comment, header)
                    if creative_commons:
                        header_adj += '\n{0} MAGIC %md\n{0} MAGIC {1}'.format(
                            self.base_comment, _cc_license
                        )
                        is_first = False

                output.write(header_adj + '\n')

                added_run = False
                for (part, code, labels, content) in commands:
                    if parts:
                        if part > i:  # don't show later parts
                            break
                        elif part == i - 1 and not added_run:
                            # add %run command before new part
                            student = NotebookGenerator.user_to_extension[
                                NotebookUser.STUDENT
                            ]
                            instr = NotebookGenerator.user_to_extension[
                                NotebookUser.INSTRUCTOR
                            ]
                            solution_file = file_out.replace(student, instr)

                            solution_file = os.path.splitext(solution_file)[0]
                            solution_file = solution_file.replace(
                                part_string, part_base.format(i)
                            )
                            # Now that Databricks supports relative run paths,
                            # this is no longer necessary.
                            #db_location = ('/Users/<EMAIL>/Labs/' +
                            #               solution_file.replace('\\', '/'))
                            # db_location = db_location.replace('build_mp/', '')
                            db_location = os.path.basename(solution_file.replace('\\', '/'))
                            runCommand = '{0} MAGIC %run {1}'.format(
                                self.base_comment, db_location
                            )
                            is_first = self._write_command(
                                output, command_cell, [runCommand, ''], is_first
{'fraction_non_alphanumeric': 0.07848934883408373, 'fraction_numerical': 0.0025535918285061487, 'mean_word_length': 1.6955261727947835, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 21, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7233062', 'n_tokens_mistral': 4123, 'n_tokens_neox': 3943, 'n_words': 898}