<commit_before>from __future__ import print_function
import abc
import types
import os
import random
from functools import wraps, partial
from itertools import chain
from collections import Container, namedtuple
from six import add_metaclass
import warnings
import math

from gbdxtools.ipe.io import to_geotiff
from gbdxtools.ipe.util import RatPolyTransform, AffineTransform, pad_safe_positive, pad_safe_negative, IPE_TO_DTYPE, preview

from shapely import ops, wkt
from shapely.geometry import box, shape, mapping, asShape
from shapely.geometry.base import BaseGeometry
try:
    from rio_hist.match import histogram_match
except ImportError:
    pass
import mercantile

import skimage.transform as tf

import pyproj
import dask
from dask import sharedict, optimize
from dask.delayed import delayed
import dask.array as da
from dask.base import is_dask_collection
from dask.rewrite import RewriteRule, RuleSet
import numpy as np

import operator
from toolz import valmap

from affine import Affine

try:
    from matplotlib import pyplot as plt
    has_pyplot = True
except:
    has_pyplot = False

try:
    xrange
except NameError:
    xrange = range

threads = int(os.environ.get('GBDX_THREADS', 8))
threaded_get = partial(dask.threaded.get, num_workers=threads)

class DaskMeta(namedtuple("DaskMeta", ["dask", "name", "chunks", "dtype", "shape"])):
    __slots__ = ()
    @classmethod
    def from_darray(cls, darr, new=tuple.__new__, len=len):
        dsk, _ = optimize.cull(darr.dask, darr.__dask_keys__())
        itr = [dsk, darr.name, darr.chunks, darr.dtype, darr.shape]
        return cls._make(itr)

    @property
    def values(self):
        return self._asdict().values()

    def infect(self, target):
        assert isinstance(target, da.Array), "DaskMeta can only be attached to Dask Arrays"
        assert len(target.shape) in [2, 3], "target must be a dask array with 2 or 3 dimensions"
        target.__dict__["__daskmeta__"] = property(lambda s: self, DaskImage.__set_daskmeta__)
        return target

class DaskImage(da.Array):
    """
    A DaskImage is a 2 or 3 dimension dask array that contains implements the `__daskmeta__` interface.
    """
    def __new__(cls, dm, **kwargs):
        if isinstance(dm, da.Array):
            dm = DaskMeta.from_darray(dm)
        elif isinstance(dm, dict):
            dm = DaskMeta(**dm)
        elif isinstance(dm, DaskMeta):
            pass
        elif dm.__class__.__name__ == "Op":
            itr = [dm.dask, dm.name, dm.chunks, dm.dtype, dm.shape]
            dm = DaskMeta._make(itr)
        else:
            raise ValueError("{} must be initialized with a DaskMeta, a dask array, or a dict with DaskMeta fields".format(cls.__name__))
        self = da.Array.__new__(cls, *dm.values)
        self.__geo_transform__ = kwargs.get("__geo_transform__")
        self.__geo_interface__ = kwargs.get("__geo_interface__")
        return self

    @property
    def __daskmeta__(self):
        return DaskMeta(self)

    @classmethod
    def create(cls, dm, **kwargs):
        """
        Given a dask meta object, construct a dask array, attach dask meta object.
        """
        assert isinstance(dm, DaskMeta), "argument must be an instance of a DaskMeta subclass"
        with dask.set_options(array_plugins=[dm.infect]):

            return da.Array.__new__(cls, dm.dask, dm.name, dm.chunks, dm.dtype, dm.shape)

    def read(self, bands=None, **kwargs):
        """
        Reads data from a dask array and returns the computed ndarray matching the given bands

        kwargs:<commit_msg>checking for graphmeta. these kinds of things might be refactored out in part 2<commit_after>from __future__ import print_function
import abc
import types
import os
import random
from functools import wraps, partial
from itertools import chain
from collections import Container, namedtuple
from six import add_metaclass
import warnings
import math

from gbdxtools.ipe.io import to_geotiff
from gbdxtools.ipe.util import RatPolyTransform, AffineTransform, pad_safe_positive, pad_safe_negative, IPE_TO_DTYPE, preview

from shapely import ops, wkt
from shapely.geometry import box, shape, mapping, asShape
from shapely.geometry.base import BaseGeometry
try:
    from rio_hist.match import histogram_match
except ImportError:
    pass
import mercantile

import skimage.transform as tf

import pyproj
import dask
from dask import sharedict, optimize
from dask.delayed import delayed
import dask.array as da
from dask.base import is_dask_collection
from dask.rewrite import RewriteRule, RuleSet
import numpy as np

import operator
from toolz import valmap

from affine import Affine

try:
    from matplotlib import pyplot as plt
    has_pyplot = True
except:
    has_pyplot = False

try:
    xrange
except NameError:
    xrange = range

threads = int(os.environ.get('GBDX_THREADS', 8))
threaded_get = partial(dask.threaded.get, num_workers=threads)

class DaskMeta(namedtuple("DaskMeta", ["dask", "name", "chunks", "dtype", "shape"])):
    __slots__ = ()
    @classmethod
    def from_darray(cls, darr, new=tuple.__new__, len=len):
        dsk, _ = optimize.cull(darr.dask, darr.__dask_keys__())
        itr = [dsk, darr.name, darr.chunks, darr.dtype, darr.shape]
        return cls._make(itr)

    @property
    def values(self):
        return self._asdict().values()

    def infect(self, target):
        assert isinstance(target, da.Array), "DaskMeta can only be attached to Dask Arrays"
        assert len(target.shape) in [2, 3], "target must be a dask array with 2 or 3 dimensions"
        target.__dict__["__daskmeta__"] = property(lambda s: self, DaskImage.__set_daskmeta__)
        return target

class DaskImage(da.Array):
    """
    A DaskImage is a 2 or 3 dimension dask array that contains implements the `__daskmeta__` interface.
    """
    def __new__(cls, dm, **kwargs):
        if isinstance(dm, da.Array):
            dm = DaskMeta.from_darray(dm)
        elif isinstance(dm, dict):
            dm = DaskMeta(**dm)
        elif isinstance(dm, DaskMeta):
            pass
        elif dm.__class__.__name__ in ("Op", "GraphMeta"):
            itr = [dm.dask, dm.name, dm.chunks, dm.dtype, dm.shape]
            dm = DaskMeta._make(itr)
        else:
            raise ValueError("{} must be initialized with a DaskMeta, a dask array, or a dict with DaskMeta fields".format(cls.__name__))
        self = da.Array.__new__(cls, *dm.values)
        self.__geo_transform__ = kwargs.get("__geo_transform__")
        self.__geo_interface__ = kwargs.get("__geo_interface__")
        return self

    @property
    def __daskmeta__(self):
        return DaskMeta(self)

    @classmethod
    def create(cls, dm, **kwargs):
        """
        Given a dask meta object, construct a dask array, attach dask meta object.
        """
        assert isinstance(dm, DaskMeta), "argument must be an instance of a DaskMeta subclass"
        with dask.set_options(array_plugins=[dm.infect]):

            return da.Array.__new__(cls, dm.dask, dm.name, dm.chunks, dm.dtype, dm.shape)

    def read(self, bands=None, **kwargs):
        """
        Reads data from a dask array and returns the computed ndarray matching the given bands

        kwargs:
{'fraction_non_alphanumeric': 0.08930800166412425, 'fraction_numerical': 0.0020801553182637636, 'mean_word_length': 3.3031026252983295, 'pattern_counts': {'":': 1, '<': 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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1844290', 'n_tokens_mistral': 2276, 'n_tokens_neox': 2222, 'n_words': 706}