<commit_before>from collections import Iterator

import os

from os import mkdir
from os.path import exists, isdir

try:
    import cPickle as pickle
except ImportError:
    import pickle

import shutil
import tempfile

from functools import partial

import blosc
import bloscpack

import numpy as np
import pandas as pd

from pandas import msgpack


def escape(text):
    return str(text)


def _safe_mkdir(path):
    if not exists(path):
        mkdir(path)


class Castra(object):
    def __init__(self, path=None, template=None, categories=None):
        # check if we should create a random path
        if path is None:
            self.path = tempfile.mkdtemp(prefix='castra-')
            self._explicitly_given_path = False
        else:
            self.path = path
            self._explicitly_given_path = True

        # check if the given path exists already and create it if it doesn't
        if not exists(self.path):
            mkdir(self.path)
        # raise an Exception if it isn't a directory
        elif not isdir(self.path):
            raise ValueError("'path': %s must be a directory")

        # either we have a meta directory
        if exists(self.dirname('meta')) and isdir(self.dirname('meta')):
            if template is not None:
                raise ValueError(
                    "'template' must be 'None' when opening a Castra")
            self.load_meta()
            self.load_partitions()
            self.load_categories()

        # or we don't, in which case we need a template
        elif template is not None:
            self.columns, self.dtypes, self.index_dtype = \
                list(template.columns), template.dtypes, template.index.dtype
            self.axis_names = [template.index.name, template.columns.name]
            self.partitions = pd.Series([], dtype='O',
                                        index=template.index.__class__([]))
            self.minimum = None
            if isinstance(categories, (list, tuple)):
                self.categories = dict((col, []) for col in categories)
            elif categories is True:
                self.categories = dict((col, [])
                                       for col in template.columns
                                       if template.dtypes[col] == 'object')
            else:
                self.categories = dict()


            mkdir(self.dirname('meta'))
            mkdir(self.dirname('meta', 'categories'))
            self.flush_meta()
            self.save_partitions()
        else:<commit_msg>Make sure our categories are properly typed in Castra<commit_after>from collections import Iterator

import os

from os import mkdir
from os.path import exists, isdir

try:
    import cPickle as pickle
except ImportError:
    import pickle

import shutil
import tempfile

from functools import partial

import blosc
import bloscpack

import numpy as np
import pandas as pd

from pandas import msgpack


def escape(text):
    return str(text)


def _safe_mkdir(path):
    if not exists(path):
        mkdir(path)


class Castra(object):
    def __init__(self, path=None, template=None, categories=None):
        # check if we should create a random path
        if path is None:
            self.path = tempfile.mkdtemp(prefix='castra-')
            self._explicitly_given_path = False
        else:
            self.path = path
            self._explicitly_given_path = True

        # check if the given path exists already and create it if it doesn't
        if not exists(self.path):
            mkdir(self.path)
        # raise an Exception if it isn't a directory
        elif not isdir(self.path):
            raise ValueError("'path': %s must be a directory")

        # either we have a meta directory
        if exists(self.dirname('meta')) and isdir(self.dirname('meta')):
            if template is not None:
                raise ValueError(
                    "'template' must be 'None' when opening a Castra")
            self.load_meta()
            self.load_partitions()
            self.load_categories()

        # or we don't, in which case we need a template
        elif template is not None:
            self.columns, self.dtypes, self.index_dtype = \
                list(template.columns), template.dtypes, template.index.dtype
            self.axis_names = [template.index.name, template.columns.name]
            self.partitions = pd.Series([], dtype='O',
                                        index=template.index.__class__([]))
            self.minimum = None
            if isinstance(categories, (list, tuple)):
                self.categories = dict((col, []) for col in categories)
            elif categories is True:
                self.categories = dict((col, [])
                                       for col in template.columns
                                       if template.dtypes[col] == 'object')
            else:
                self.categories = dict()

            if self.categories:
                if set(self.categories) - set(template.dtypes.index.values):
                    raise ValueError('passed in categories %s are not all '
                                     'contained in template dataframe columns '
                                     '%s' %
                                     (set(self.categories),
                                      set(template.dtypes.index.values)))
                for c in self.categories:
                    self.dtypes[c] = pd.core.categorical.CategoricalDtype()

            mkdir(self.dirname('meta'))
            mkdir(self.dirname('meta', 'categories'))
            self.flush_meta()
            self.save_partitions()
        else:
{'fraction_non_alphanumeric': 0.08502024291497975, 'fraction_numerical': 0.0, 'mean_word_length': 1.7595920349684313, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3252568', 'n_tokens_mistral': 1481, 'n_tokens_neox': 1420, 'n_words': 456}