text:
import ast
from ast import ClassDef, FunctionDef
import builtins
from pathlib import Path
from .agenda_util import pprint_agenda
from .deprecations import pprint_def_names
from .import_util import get_imported_name_sources, annotate_imports
from .def_path_util import (
    UntypedPathStr,
    FuncDefPathStr,
    InnerFuncDefPathStr,
    ClassDefPathStr,
    HigherOrderClassDefPathStr,
    InnerClassDefPathStr,
    MethodDefPathStr,
)
from .def_helpers import _find_node, _find_cls, _find_def
from sys import stderr
from itertools import chain
from functools import reduce
from enum import Enum

__all__ = [
    "retrieve_ast_agenda",
    "process_ast",
    "find_assigned_args",
    "set_extradef_names",
    "set_nondef_names",
    "get_def_names",
    "parse_mv_funcs",
]

def retrieve_ast_agenda(linkfile, transfers=None):
    """
    Build and parse the Abstract Syntax Tree (AST) of a Python file, and either return
    a report of what changes would be required to move the mvdefs subset of all
    function definitions out of it, or a report of the imports and funcdefs in general
    if no `linkfile.mvdefs` is provided (taken to indicate that the file is the target funcdefs
    are moving to), or make changes to the file (either newly creating one if no such
    file exists, or editing in place according to the reported import statement
    differences).

    If the Python file `linkfile.path` doesn't exist (which can be checked directly as it's a
    Path object), it's being newly created by the move and obviously no report can
    be made on it: it has no funcdefs and no import statements, so all the ones being
    moved will be newly created.

    mvdefs should be given if the file is the source of moved functions, and left
    empty (default: `None` which --> `[]`) if the file is the destination to move them to.

    If `linkfile.report` is True, returns a string describing the changes
    to be made (if False, nothing is returned).

    If backup is True, files will be changed in place by calling `mvdef.backup.backup`
    (obviously, be careful switching this setting off if report is True, as any
    changes made cannot be restored afterwards from this backup file).
    """
    if linkfile.is_extant:
        with open(linkfile.path, "r") as f:
            fc = f.read()
            trunk = ast.parse(fc).body

        # print("Next running process_ast from retrieve_ast_agenda")
        linkfile.process_ast(trunk, transfers)  # sets linkfile.edits
    elif type(linkfile).__name__ == "DstFile":
        # An `isinstance` call would require a circular import, hence the __name__ check
        #
        # Not extant so file doesn't exist (cannot produce a parsed AST) however the
        # linkfile is the destination (no `mvdefs` to remove), so return None.
        # This will be picked up by the assert in SrcFile.validate_edits
        # (but skipped for DstFile.validate_edits)
        assert linkfile.mvdefs is None, "Unexpected mvdefs list for non-extant DstFile"
        linkfile.edits = None
    else:
        msg = f"Can't move {linkfile.mvdefs=} from {linkfile.path=} – it doesn't exist!"
        raise ValueError(msg)
    # print("Finished processing in retrieve_ast_agenda")

class EditAgenda(dict):
    def __init__(self):
        super().__init__({c: [] for c in self.categories})

    categories = ["move", "keep", "copy", "lose", "take", "echo", "stay"]

    def add_entry(self, category, key_val_pair):
        if category not in self.categories:
            raise ValueError(f"{category} is not a valid agenda category")
        entry = EditItem(*key_val_pair)
        self.get(category).append(entry)

    def remove_entry(self, category, entry_value):
        index_key = [next(iter(x.values())) for x in self.get(category)].index(
            entry_value
        )
        del self.get(category)[index_key]

    def add_imports(self, imports, category, names):
        for i in imports:
            assert i in set().union(*names.values()), f"{i} not found"
            i_dict = next(v.get(i) for k, v in names.items() if i in v)
            self.add_entry(category=category, key_val_pair=(i, i_dict))

class EditItem(dict):
    def __init__(self, key, value):
        super().__init__({key: value})

def process_ast(linkfile, trunk, transfers=None):
    """
    Handle the hand-off to dedicated functions to go from the mvdefs of functions
    to move, first deriving lists of imported names which belong to the mvdefs and
    the non-mvdefs functions (using `parse_mv_funcs`), then constructing an
    'edit agenda' (using `process_ast`) which describes [and optionally
    reports] the changes to be made at the file level, in terms of move/keep/copy
    operations on individual import statements between the source and destination
    Python files.

      mvdefs:     List of functions to be moved
      trunk:      Tree body of the file's AST, which will be separated into
                  function definitions, import statements, and anything else.
      transfers:  List of transfers already determined to be made from the src
                  to the dst file (from the first call to ast_parse)
      report:     Whether to print a report during the program (default: True)

    -------------------------------------------------------------------------------

    First, given the lists of mvdef names (linkfile.mvdef_names) and non-mvdef names
    (linkfile.nonmvdef_names), construct the subsets:

      mv_imps:  imported names used by the functions to move (only in mvdef_names),
      nm_imps:  imported names used by the functions not to move (only in
                nonmvdef_names),
      mu_imps:  imported names used by both the functions to move and the
                functions not to move (in both mvdef_names and nonmvdef_names)

    Potentially 'as a dry run' (if this is being called by process_ast and its
    parameter edit is False), report how to remove the import statements or statement
    sections which import mv_inames, do nothing to the import statements which import
    nonmv_inames, and copy the import statements which import mutual_inames (as both
    src and dst need them).

    Additionally, accept 'transfers' from a previously determined edit agenda,
    so as to "take" the "move" names, and "echo" the "copy" names (i.e. when
    receiving names marked by "move" and "copy", distinguish them to indicate
    they are being received by transfer [from src⇒dst file], for clarity).

    For clarity, note that this function does **not** edit anything itself, it just
    describes how it would be possible to carry out the required edits at the level
    of Python file changes. As such it is computed even on a 'dry run'.
    """
    linkfile.parse_mv_funcs(trunk)  # sets ast_funcs, {mv,nonmv,extra,un}def_names
    imported_names = get_imported_name_sources(trunk, report=linkfile.report)
    if linkfile.report:
        print(f"• Determining edit agenda for {linkfile.path.name}:", file=stderr)
    linkfile.edits = EditAgenda()
    linkfile.imp_def_subsets()  # sets mv_imports, nonmv_imports, mutual_imports
    # Iterate over each imported name, i, in the subset of import names to move

    linkfile.edits.add_imports(linkfile.mv_imports, "move", linkfile.mvdef_names)
    linkfile.edits.add_imports(linkfile.nonmv_imports, "keep", linkfile.nonmvdef_names)
    linkfile.edits.add_imports(linkfile.mutual_imports, "copy", linkfile.mvdef_names)
    for i in linkfile.undef_names:
        i_dict = linkfile.undef_names.get(i)
        linkfile.edits.add_entry(category="lose", key_val_pair=(i, i_dict))
    if not transfers:
        # Returning without transfers if None (would also catch empty dict `{}`)
        if linkfile.report:
            pprint_agenda(linkfile.edits)
        return
    # elif linkfile.report:
    #    if len(linkfile.edits.get("lose")) > 0:
    #        print("• Resolving edit agenda conflicts:")
    # i is 'ready made' from a previous call to ast_parse, and just needs reporting
    for i in transfers.get("take"):
        k, i_dict = next(iter((i.items())))
        linkfile.edits.add_entry(category="take", key_val_pair=(k, i_dict))
    for i in transfers.get("echo"):
        k, i_dict = next(iter((i.items())))
        linkfile.edits.add_entry(category="echo", key_val_pair=(k, i_dict))
    # Resolve agenda conflicts: if any imports marked 'lose' are cancelled out
    # by any identically named imports marked 'take' or 'echo', change to 'stay'
    for i in linkfile.edits.get("lose"):
        k, i_dict = next(iter((i.items())))
        imp_src = i_dict.get("import")
        if k in [list(x)[0] for x in linkfile.edits.get("take")]:
            t_i_dict = next(x.get(k) for x in linkfile.edits.get("take") if k in x)
            take_imp_src = t_i_dict.get("import")
            if imp_src != take_imp_src:
                continue
            # Deduplicate 'lose'/'take' k: replace both with 'stay'
            linkfile.edits.add_entry(category="stay", key_val_pair=(k, i_dict))
            linkfile.edits.remove_entry(category="lose", entry_value=i_dict)
            linkfile.edits.remove_entry(category="take", entry_value=t_i_dict)
        elif k in [list(x)[0] for x in linkfile.edits.get("echo")]:
            e_i_dict = next(x.get(k) for x in linkfile.edits.get("echo") if k in x)
            echo_imp_src = e_i_dict.get("import")
            if imp_src != echo_imp_src:
                continue
            # Deduplicate 'lose'/'echo' k: replace both with 'stay'
            linkfile.edits.add_entry(category="stay", key_val_pair=(k, i_dict))
            linkfile.edits.remove_entry(category="lose", entry_value=i_dict)
            linkfile.edits.remove_entry(category="echo", entry_value=e_i_dict)
    # Resolve agenda conflicts: if any imports marked 'take' or 'echo' are cancelled
    # out by any identically named imports already present, change to 'stay'
    for i in linkfile.edits.get("take"):
        k, i_dict = next(iter((i.items())))
        take_imp_src = i_dict.get("import")
        if k in imported_names:
            # Check the import source and asnames match
            imp_src = imported_names.get(k)[0]
            if imp_src != take_imp_src:
                # This means that the same name is being used by a different function
                raise ValueError(
                    f"Cannot move imported name '{k}', it is already "
                    + f"in use in {linkfile.path.name} ({take_imp_src} clashes with {imp_src})"
                )
                # (N.B. could rename automatically as future feature)
            # Otherwise there is simply a duplicate import statement, so the demand
            # to 'take' the imported name is already fulfilled.
            # Replace unnecessary 'take' with 'stay'
            linkfile.edits.add_entry(category="stay", key_val_pair=(k, i_dict))
            linkfile.edits.remove_entry(category="take", entry_value=i_dict)
    for i in linkfile.edits.get("echo"):
        k, i_dict = next(iter((i.items())))
        echo_imp_src = i_dict.get("import")
        if k in imported_names:
            # Check the import source and asnames match
            imp_src = imported_names.get(k)[0]
            if imp_src != echo_imp_src:
                # This means that the same name is being used by a different function
                raise ValueError(
                    f"Cannot move imported name '{k}', it is already "
                    + f"in use in {linkfile.path.name} ({echo_imp_src} clashes with {imp_src})"
                )
                # (N.B. could rename automatically as future feature)
            # Otherwise there is simply a duplicate import statement, so the demand
            # to 'echo' the imported name is already fulfilled.
            # Replace unnecessary 'echo' with 'stay'
            linkfile.edits.add_entry(category="stay", key_val_pair=(k, i_dict))
            linkfile.edits.remove_entry(category="echo", entry_value=i_dict)
    if linkfile.report:
        pprint_agenda(linkfile.edits)
    return

def find_assigned_args(fd):
    """
    Produce a list of the names in a function definition `fd` which are created
    by assignment operations (as identified via the function definition's AST).
    """
    args_indiv = []  # Arguments assigned individually, e.g. x = 1
    args_multi = []  # Arguments assigned from a tuple, e.g. x, y = (1,2)
    for a in ast.walk(fd):
        if type(a) is ast.Assign:
            # Handle explicit assignments from use of the equals symbol
            assert len(a.targets) == 1, "Expected 1 target per ast.Assign"
            if type(a.targets[0]) is ast.Name:
                args_indiv.append(a.targets[0].id)
            elif type(a.targets[0]) is ast.Tuple:
                args_multi.extend([x.id for x in a.targets[0].elts])
        elif type(a) is ast.For:
            # Handle implicit assignments (ctx = Store) within for loops
            if type(a.target) is ast.Name:
                args_indiv.append(a.target.id)
            elif type(a.target) is ast.Tuple:
                for x in a.target.elts:
                    assert type(x) in [ast.Name, ast.Tuple], f"Unexpected target.elts"
                    if type(x) is ast.Name:
                        args_multi.append(x.id)
                    else:  # type(x) is ast.Tuple
                        for y in x.elts:
                            assert type(y) is ast.Name, f"Unexpected target.elts tuple"
                            args_multi.append(y.id)
            else:
                raise ValueError(f"{a.target} lacks the expected ast.Name statements")
        elif type(a) in (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp):
            # Handle implicit assignments within comprehensions
            for g in a.generators:
                assert type(g) is ast.comprehension, f"{a} not ast.comprehension type"
                if type(g.target) is ast.Name:
                    args_indiv.append(g.target.id)
                elif type(g.target) is ast.Tuple:
                    args_multi.extend([x.id for x in g.target.elts])
                else:
                    raise ValueError(f"{g.target} lacks expected ast.Name statements")
        elif type(a) is ast.Lambda:
            if len(a.args.args) > 1:
                args_multi.extend([r.arg for r in a.args.args])
            else:
                assert len(a.args.args) == 1, "A lambda can't assign no names (can it?)"
                args_indiv.append(a.args.args[0].arg)
        elif type(a) is ast.NamedExpr:
            # This is the walrus operator `:=` added in Python 3.8
            assert type(a.target) is ast.Name, f"Expected a name for {a.target}"
            args_indiv.append(a.target.id)
            # I haven't seen a multiple assignment from a walrus operator, asked SO:
            # https://stackoverflow.com/q/59567172/2668831
    assigned_args = args_indiv + args_multi
    return assigned_args

def set_extradef_names(linkfile, extra_nodes):
    """
    Return the names used in the AST trunk nodes which are outside of both function
    definitions and import statements, so as to distinguish the unused names from
    those which are just used outside of function definitions.
    """
    linkfile.extradef_names = set()
    for node in extra_nodes:
        node_names = [x.id for x in ast.walk(node) if type(x) is ast.Name]
        for n in node_names:
            linkfile.extradef_names.add(n)
    return

def set_nondef_names(linkfile, unused, import_annos):
    report = linkfile.report
    imp_name_lines, imp_name_dicts = import_annos
    linkfile.nondef_names = NameDict(unused)  # dict keyed by the unused names
    if unknowns := [n for n in unused if n not in imp_name_lines]:
        raise ValueError(f"These names could not be sourced: {unknowns}")
    # mv_imp_refs is the subset of imp_name_lines for movable funcdef names
    # These refs will lead to import statements being copied and/or moved
    uu_imp_refs = {n: imp_name_lines.get(n) for n in unused}
    for k in uu_imp_refs:  # iterate over the unused imported names/asnames
        n = uu_imp_refs.get(k).get("n")
        d = imp_name_dicts[n]
        n_i = next(i for i, x in enumerate(d.values()) if x == k)
        assert n_i >= 0, f"Movable name {k} not found in import name dict"
        # Store index in case of multiple imports per import statement line
        uu_imp_refs.get(k)["n_i"] = n_i
        new_entry = {
            "n": uu_imp_refs.get(k).get("n"),
            "n_i": n_i,
            "line": uu_imp_refs.get(k).get("line"),
            "import": [*imp_name_dicts[n]][n_i],
        }
        linkfile.nondef_names.add_entry(k, new_entry)
    return

class NameDict(dict):
    def __init__(self, def_list):
        super().__init__({f: {} for f in def_list})

    def add_entry(self, def_name, def_name_info_dict):
        if self.get(def_name):
            raise ValueError(f"{self.get(def_name)=} already exists!")
        self[def_name] = def_name_info_dict

    def add_subentry(self, def_name, import_name, import_name_info_dict):
        if self.get(def_name).get(import_name):
            raise ValueError(f"{self.get(def_name).get(import_name)=} already exists!")
        self[def_name][import_name] = import_name_info_dict

class NameEntryDict(dict):
    # can only be used for `get_def_names`, the others are not consistent enough
    # even though they are dicts with the same keys...
    def __init__(self, names, values=None, sort=False):
        if values:
            if sort:
                raise NotImplementedError("Sorting not coded for key-val pairs")
            super().__init__(dict(zip(names, values)))
        else:
            if type(names) is dict:
                if sort:
                    raise NotImplementedError("Sorting not coded for dicts")
                super().__init__(names)  # simples
                return
            elif sort:
                names = sorted(names)
            super().__init__({n: {} for n in names})

def get_base_type_name(type_name):
    """
    Take a type name from the IntraDefTypeEnum names and return the name of the base
    type, e.g. input the MethodDef type and get out the name "Func", indicating that a
    method is a function
    """
    msg = f"{type_name} is not an inner def type name (i.e. a name in IntraDefTypeEnum)"
    if type_name in IntraDefTypeEnum._member_names_:
        def_bases = DefTypeEnum._value2member_map_
        type_class = IntraDefTypeEnum[type_name].value
        base_type = next(b for b in type_class.__bases__ if b in def_bases)
        base_type_name = DefTypeEnum(base_type).name
    else:
        if type_name in DefTypeEnum._member_names_:
            base_type_name = type_name  # already a base type name ("Class" or "Func")
        else:
            raise ValueError(msg)
    return base_type_name

def has_clsdef_base(path_part, intradef=True):
    next_def_type_name = path_part.part_type  # e.g. "Method"
    next_def_base_type_name = get_base_type_name(next_def_type_name)  # Method --> Func
    def_is_a_cls = next_def_base_type_name == "Class"
    return def_is_a_cls

class PathGetterMixin:
    """
    Used in all `Def` classes (`FuncDef`, `ClsDef` and all subclasses) to provide
    a generic way to retrieve the inner funcdef/clsdef so as to retrieve a path
    to a leaf funcdef/clsdef.

    (I don't think this even needs to be mixed in but it's handy to access if it is)
    """

    def retrieve_def(self, path_part):
        """
        Attempt to retrieve the next funcdef/clsdef indicated by `path_part`
        (an instance one of the `PathPart` classes iterated over via `reduce`
        in `check_against_linkedfile` from one of the `Path` classes).
        """
        # can use either IntraDefPathTypeEnum/IntraDefTypeEnum (latter more consistent)
        if path_part.part_type not in IntraDefTypeEnum._member_map_:
            if path_part in DefTypeEnum._member_map_:
                # handle global defs in check_against_linkedfile to get an initial def
                msg = f"{path_part=} is not supposed to be a top-level def type"
            else:
                msg = f"{path_part=} is not a recognised def type"
            raise TypeError(f"{msg} (path_part.part_type=)")
        def_is_a_cls = has_clsdef_base(path_part)
        intras = self.intra_classes if def_is_a_cls else self.intra_funcs
        retrieved_def = _find_node(intras, path_part)
        return retrieved_def

    def retrieve_def_from_body(self, path_part):
        """
        Attempt to retrieve the next funcdef/clsdef indicated by `path_part`
        (an instance one of the `PathPart` classes iterated over via `reduce`
        in `check_against_linkedfile` from one of the `Path` classes). Use the
        body rather than `intra_classes` and `intra_funcs`
        """
        # can use either IntraDefPathTypeEnum/IntraDefTypeEnum (latter more consistent)
        if path_part.part_type not in IntraDefTypeEnum._member_map_:
            if path_part in DefTypeEnum._member_map_:
                # handle global defs in check_against_linkedfile to get an initial def
                msg = f"{path_part=} is not supposed to be a top-level def type"
            else:
                msg = f"{path_part=} is not a recognised def type"
            raise TypeError(f"{msg} (path_part.part_type=)")
        def_is_a_cls = has_clsdef_base(path_part)
        finder = _find_cls if def_is_a_cls else _find_def
        retrieved_def = finder(self, path_part)
        return retrieved_def

class LinkFileCheckerMixin:
    def check_against_linkedfile(self, linkfile):
        filename = linkfile.path.name
        if (
            linkfile.classes_only
        ):  # use pre-extracted IDs rather than do fresh listcomps
            global_cls_names, global_fun_names = linkfile.sel_ids, linkfile.nosel_ids
        else:
            global_fun_names, global_cls_names = linkfile.sel_ids, linkfile.nosel_ids
        if (self.root_type == "Class" and self.root_name not in global_cls_names) or (
            self.root_type == "Func" and self.root_name not in global_fun_names
        ):
            msg = f"{filename} does not contain root {self.root_type} {self.root_name}"
            raise NameError(msg)
        else:
            global_classes, global_funcs = linkfile.ast_classes, linkfile.ast_funcs
            poss_roots = global_classes if self.root_type == "Class" else global_funcs
            initial_def = _find_node(poss_roots, self.root_name)
            remaining_parts = self.parts[1:]
            if remaining_parts:
                try:
                    retrieved_def = reduce(
                        PathGetterMixin.retrieve_def, remaining_parts, initial_def
                    )
                except Exception as e:
                    msg = f"{filename} does not contain {self.string} (447raised {e})"
                    raise NameError(msg)
            else:
                retrieved_def = initial_def
            return retrieved_def

    def check_against_defs(self, global_def_nodes):
        assert global_def_nodes, "check_against_defs was called with an empty list"
        global_def_names = [d.name for d in global_def_nodes]
        is_cls = lambda c: isinstance(c, ClassDef)
        is_fun = lambda f: isinstance(f, FunctionDef)
        cls_check = self.root_type == "Class" and all(map(is_cls, global_def_nodes))
        fun_check = self.root_type == "Func" and all(map(is_fun, global_def_nodes))
        msg = f"Defs passed were not of type {self.root_type}"
        assert cls_check or fun_check, msg
        initial_def = _find_node(global_def_nodes, self.root_name)
        msg = f"{global_def_nodes} does not contain root {self.root_type} {self.root_name}"
        assert initial_def is not None, msg
        remaining_parts = self.parts[1:]
        if remaining_parts:
            try:
                retrieved_def = reduce(
                    PathGetterMixin.retrieve_def_from_body, remaining_parts, initial_def
                )
            except Exception as e:
                msg = f"Failed to retrieve {self.string} (472raised {e})"
                raise NameError(msg)
        else:
            retrieved_def = initial_def
        return retrieved_def

class FuncPath(FuncDefPathStr, LinkFileCheckerMixin):
    """
    A FuncDefPathStr for a top-level function definition
    """

    # fall through to FuncDefPathStr.__init__, setting .string, ._tokens and .parts
    def __init__(self, path_string, parent_type_name=None):
        if parent_type_name:
            self.parent_type_name = parent_type_name
        super().__init__(path_string)

class InnerFuncPath(InnerFuncDefPathStr, LinkFileCheckerMixin):
    """
    An InnerFuncDefPathStr which has a top level funcdef, an 'intradef' inner func
    (checked on super().__init__), and potentially one or more inner functions below
    that, which must be reachable as direct descendants of the AST at each step (i.e.
    no intervening nodes between descendant inner functions in the path when checking
    against the LinkedFile AST).
    """

    # fall through to FuncDefPathStr.__init__, setting .string, ._tokens and .parts
    def __init__(self, path_string, parent_type_name=None):
        if parent_type_name:
            self.parent_type_name = parent_type_name
        super().__init__(path_string)

class ClassPath(ClassDefPathStr, LinkFileCheckerMixin):
    """
    A ClassDefPathStr for a top-level class
    """

    # fall through to ClassDefPathStr.__init__, setting .string, ._tokens and .parts
    def __init__(self, path_string, parent_type_name=None):
        if parent_type_name:
            self.parent_type_name = parent_type_name
        super().__init__(path_string)

class InnerClassPath(InnerClassDefPathStr, LinkFileCheckerMixin):
    """
    A ClassDefPathStr for a class within another class.
    """

    # fall through to ClassDefPathStr.__init__, setting .string, ._tokens and .parts
    def __init__(self, path_string, parent_type_name=None):
        if parent_type_name:
            self.parent_type_name = parent_type_name
        super().__init__(path_string)

class HigherOrderClassPath(HigherOrderClassDefPathStr, LinkFileCheckerMixin):
    """
    A ClassDefPathStr for a class within another class.
    """

    # fall through to ClassDefPathStr.__init__, setting .string, ._tokens and .parts
    def __init__(self, path_string, parent_type_name=None):
        if parent_type_name:
            self.parent_type_name = parent_type_name
        super().__init__(path_string)

class MethodPath(MethodDefPathStr, LinkFileCheckerMixin):
    """
    A MethodDefPathStr which has a top level class, which contains a method
    (checked on super().__init__), and potentially one or more inner functions below
    that*, which must be reachable as direct descendants of the AST at each step (i.e.
    no intervening nodes between descendant inner functions in the path when checking
    against the LinkedFile AST).
    *[TODO: confirm against finished implementation if this is the case RE: inner funcs]
    """

    # fall through to FuncDefPathStr.__init__, setting .string, ._tokens and .parts
    def __init__(self, path_string, parent_type_name=None):
        if parent_type_name:
            self.parent_type_name = parent_type_name
        super().__init__(path_string)

def get_def_names(linkfile, def_list, import_annos):
    """
    Given the `def_list` list of strings (must be empty in the negative case rather
    than None, and is prepared as such from `LinkedFile.mvdefs` in `parse_mv_funcs`),
    return a dict from its keys whose entries are the names of its descendant AST nodes.
    """
    get_cls = linkfile.classes_only
    imp_name_lines, imp_name_dicts = import_annos
    def_namedict = NameDict(def_list)
    extradef_names = linkfile.extradef_names
    # ast_funcs/ast_classes parameterised as "selected" based on linkfile.classes_only
    sel_nodes, nosel_nodes = linkfile.sel_nodes, linkfile.nosel_nodes
    sel_ids, nosel_ids = linkfile.sel_ids, linkfile.nosel_ids  # nosel_ids not used?
    for m in def_list:
        m_type = "class" if get_cls else "function"
        sd_names = set()
        m_parsed = UntypedPathStr(m)  # will be remade in InnerFuncPath but it's fast
        if len(m_parsed.parts) > 1:
            root_node = m_parsed.parts[0]
            # use pen/ultimate nodes in the path (i.e. leaf and its parent)
            leaf_parent_node, leaf_node = m_parsed.parts[-2:]
            if get_cls:
                valid_leaf_types = ["InnerClass", "HigherOrderClass"]
            else:
                valid_leaf_types = ["Method", "InnerFunc"]
            msg = f"'{leaf_node.part_type}' is not a valid {m_type} part type"
            assert leaf_node.part_type in valid_leaf_types, msg
            leaf_par_type_name = getattr(
                DefTypeToParentTypeEnum, leaf_node.part_type
            ).value
            m_path_type = getattr(IntraDefPathTypeEnum, leaf_node.part_type).value
            m_path = m_path_type(m, parent_type_name=leaf_par_type_name)
            # retrieve {func|cls}def from AST
            sel_def = m_path.check_against_linkedfile(linkfile)
            sel_ids = (
                sel_def.all_ns_cd_ids
                if get_cls
                else sel_def.all_ns_fd_ids
                # TODO make this a property on the type (which class though?)
            )  # inner {func|cls}def IDs, includes global def namespace
        elif m in sel_ids:
            sel_def = sel_nodes[sel_ids.index(m)]
        else:
            raise NameError(f"No {m_type} '{m}' is defined")
        # def params (sel_def may be intra)
        sd_params = [] if get_cls else [a.arg for a in sel_def.args.args]
        assigned = find_assigned_args(sel_def)
        for ast_statement in sel_def.body:
            exc = dir(builtins) + sel_ids + sd_params + assigned + [*extradef_names]
            for node in ast.walk(ast_statement):
                if type(node) == ast.Name:
                    n_id = node.id
                    if n_id not in exc:
                        sd_names.add(n_id)
        def_namedict[m] = NameEntryDict(sd_names, sort=True)
        # All names successfully found and can finish if remaining names are
        # in the set of funcdef names, comparing them tothe import statements
        unknowns = [n for n in sd_names if n not in imp_name_lines]
        if unknowns:
            raise ValueError(f"These names could not be sourced: {unknowns}")
        # mv_imp_refs is the subset of imp_name_lines for movable funcdef names
        # These refs will lead to import statements being copied and/or moved
        mv_imp_refs = dict([(n, imp_name_lines.get(n)) for n in sd_names])
        update_def_names_from_imports(m, mv_imp_refs, imp_name_dicts, def_namedict)
    return def_namedict

def update_def_names_from_imports(def_name, mv_imp_refs, imp_name_dicts, def_names):
    # TODO inspect def_names before/after this routine runs and add an informative
    # docstring/change its name to something better (named in midst of a refactor)
    for k in mv_imp_refs:
        n = mv_imp_refs.get(k).get("n")
        d = imp_name_dicts[n]
        n_i = next(list(d.keys()).index(x) for x in d if d[x] == k)
        assert n_i >= 0, f"Movable name {k} not found in import name dict"
        # Store index in case of multiple imports per import statement line
        mv_imp_refs.get(k)["n_i"] = n_i
        n = mv_imp_refs.get(k).get("n")
        new_entry = NameEntryDict(
            {
                "n": n,
                "n_i": n_i,
                "line": mv_imp_refs.get(k).get("line"),
                "import": list(imp_name_dicts[n].keys())[n_i],
            }
        )
        def_names.add_subentry(def_name, k, new_entry)  # mutate in place

def _list_names_of_type(nodes, of_type):
    return [n.name for n in nodes if isinstance(n, of_type)]

def _index_nodes_of_type(nodes, of_type):
    return [i for i, n in enumerate(nodes) if isinstance(n, of_type)]

class RecursiveIdSetterMixin:
    """
    Mixin class to provide a common interface method for both `ClsDef` and `FuncDef`
    (and by extension all inheriting subclasses).

    Note that methods have been omitted (but are trivial to obtain from the result)
    """

    def _set_up_inner_node(self, node):
        node.check_for_inner_defs()
        if self.is_inner:
            node.set_ns_cd_ids(self.all_ns_cd_ids)
            node.set_ns_fd_ids(self.all_ns_fd_ids)
        else:
            inner_cd_id_ns = [*self.global_cd_ids]
            inner_cd_id_ns.extend(c.name for c in self.intra_classes)
            inner_fd_id_ns = [*self.global_fd_ids]
            inner_fd_id_ns.extend(f.name for f in self.intra_funcs)
            node.set_ns_cd_ids(inner_cd_id_ns)
            node.set_ns_fd_ids(inner_fd_id_ns)

    def recursively_set_inner_defs(self):
        if self.has_intra_func:
            for f in self.intra_funcs:
                self._set_up_inner_node(f)
            for f in self.intra_funcs:
                f.recursively_set_inner_defs()
        if self.has_intra_cls:
            for c in self.intra_classes:
                self._set_up_inner_node(c)
            for c in self.intra_classes:
                c.recursively_set_inner_defs()

    def check_for_inner_defs(self):
        if not hasattr(self, self.intra_func_idx_attr):
            intra_func_idx = _index_nodes_of_type(self.body, ast.FunctionDef)
            setattr(self, self.intra_func_idx_attr, intra_func_idx)
        if not hasattr(self, self.intra_cls_idx_attr):
            intra_cls_idx = _index_nodes_of_type(self.body, ast.ClassDef)
            setattr(self, self.intra_cls_idx_attr, intra_cls_idx)
        self.set_intra_defs()
        if not self.is_inner:
            self.recursively_set_inner_defs()  # sets inner_funcs, all_ns_{f|c}d_ids

    def set_ns_fd_ids(self, parent_namespace_ids=None):
        self.all_ns_fd_ids = []
        if parent_namespace_ids:
            self.all_ns_fd_ids.extend(parent_namespace_ids)
        if self.has_intra_func:
            self.all_ns_fd_ids.extend(_list_names_of_type(self.body, ast.FunctionDef))

    def set_ns_cd_ids(self, parent_namespace_ids=None):
        self.all_ns_cd_ids = []
        if parent_namespace_ids:
            self.all_ns_cd_ids.extend(parent_namespace_ids)
        if self.has_intra_cls:
            self.all_ns_cd_ids.extend(_list_names_of_type(self.body, ast.ClassDef))

    @property
    def _parent_def_kwargs(self):
        return {f"parent_{'c' if isinstance(self, ClsDef) else 'f'}d": self}

    @property
    def intra_cls_type(self):
        return getattr(IntraDefTypeEnum, self.intra_cls_type_name).value

    @property
    def intra_cls_idx(self):
        return getattr(self, self.intra_cls_idx_attr)

    @property
    def intra_func_type(self):
        return getattr(IntraDefTypeEnum, self.intra_func_type_name).value

    @property
    def intra_func_idx(self):
        return getattr(self, self.intra_func_idx_attr)

    def set_intra_defs(self):
        self.set_intra_classes()
        self.set_intra_funcs()

    def set_intra_classes(self):
        self.intra_classes = [
            self.intra_cls_type(cd=self.body[i], **self._parent_def_kwargs)
            for i in self.intra_cls_idx
        ]

    def set_intra_funcs(self):
        self.intra_funcs = [
            self.intra_func_type(fd=self.body[i], **self._parent_def_kwargs)
            for i in self.intra_func_idx
        ]

class NamespaceIdSetterMixin:
    @property
    def all_ns_cd_ids(self):
        return self._all_ns_cd_ids

    @all_ns_cd_ids.setter
    def all_ns_cd_ids(self, id_list):
        self._all_ns_cd_ids = id_list

    @property
    def all_ns_fd_ids(self):
        return self._all_ns_fd_ids

    @all_ns_fd_ids.setter
    def all_ns_fd_ids(self, id_list):
        self._all_ns_fd_ids = id_list

class ClsDef(ast.ClassDef, RecursiveIdSetterMixin, PathGetterMixin):
    """
    Wrap `ast.ClassDef` to permit recursive search for methods and inner classes
    upon creation in `parse_mv_funcs`.
    """

    def __init__(
        self, clsdef, classes_only, ast_cls_ids=None, ast_fun_ids=None, is_inner=False
    ):
        super().__init__(**vars(clsdef))
        self.classes_only = classes_only
        self.global_cd_ids = ast_cls_ids
        self.global_fd_ids = ast_fun_ids
        self.is_inner = is_inner
        if not self.is_inner:
            self.check_for_inner_defs()

    def get_method(self, meth_name):
        return next(m for m in self.methods if m.name == meth_name)

    @property
    def has_method(self):
        return self.mth_idx != []

    @property
    def has_inner_cls(self):
        return self.inner_cls_idx != []

    @property
    def mth_idx(self):
        return self._mth_idx

    @mth_idx.setter
    def mth_idx(self, idx):
        self._mth_idx = idx

    def get_inner_cls(self, cls_name):
        return next(c for c in self.inner_classes if c.name == cls_name)

    @property
    def inner_cls_idx(self):
        return self._inner_cls_idx

    @inner_cls_idx.setter
    def inner_cls_idx(self, idx):
        self._inner_cls_idx = idx

    @property
    def line_range(self):
        return (self.lineno, self.end_lineno)

    def set_methods(self):
        if not self.has_method:
            self.methods = []
        self.methods = [
            MethodDef(fd=self.body[i], parent_cd=self) for i in self.mth_idx
        ]

    @property
    def methods(self):
        return self._methods

    @methods.setter
    def methods(self, mthdefs):
        self._methods = mthdefs

    @property
    def inner_classes(self):
        return self._inner_classes

    @inner_classes.setter
    def inner_classes(self, clsdefs):
        self._inner_classes = clsdefs

    # Read-only aliases used in `RecursiveIdSetterMixin.recursively_set_inner_defs`
    intra_func_idx_attr = "_mth_idx"
    intra_cls_idx_attr = "_inner_cls_idx"
    has_intra_func = has_method
    has_intra_cls = has_inner_cls
    intra_funcs = methods
    intra_classes = inner_classes
    intra_func_type_name = "Method"
    intra_cls_type_name = "InnerClass"

    @property
    def path(self):
        return self.name

class HOClsDef(ClsDef, NamespaceIdSetterMixin):
    """
    Wrap `ClsDef` (in turn wrapping `ast.ClassDef`) to store a reference to the
    parent classdef's line range on a higher order class (i.e. a class in a funcdef.
    """

    def __init__(self, cd, parent_fd):
        self.classes_only = parent_fd.classes_only
        self.parent_name = parent_fd.name
        self.parent_path = parent_fd.path
        self.parent_line_range = parent_fd.line_range
        super().__init__(
            cd, self.classes_only, is_inner=True
        )  # I moved this to the end to get it to run
        # but unsure if it's actually correct to do so or just a hotfix that'll bite me

    @property
    def path(self):
        parent_path = self.parent_path
        return f"{parent_path}:::{self.name}"

class InnerClsDef(ClsDef, NamespaceIdSetterMixin):
    """
    Wrap `ClsDef` (in turn wrapping `ast.ClassDef`) to store a reference to the
    parent classdef's line range on an inner class.
    """

    def __init__(self, cd, parent_cd):
        self.classes_only = parent_cd.classes_only
        self.parent_name = parent_cd.name
        self.parent_path = parent_cd.path
        self.parent_line_range = parent_cd.line_range
        super().__init__(
            cd, self.classes_only, is_inner=True
        )  # I moved this to the end to get it to run
        # but unsure if it's actually correct to do so or just a hotfix that'll bite me

    @property
    def path(self):
        parent_path = self.parent_path
        return f"{parent_path}::{self.name}"

class FuncDef(ast.FunctionDef, RecursiveIdSetterMixin, PathGetterMixin):
    """
    Wrap `ast.FunctionDef` to permit recursive search for inner functions upon
    creation in `parse_mv_funcs`.
    """

    def __init__(
        self, funcdef, classes_only, ast_cls_ids=None, ast_fun_ids=None, is_inner=False
    ):
        super().__init__(**vars(funcdef))
        self.classes_only = classes_only
        self.global_cd_ids = ast_cls_ids
        self.global_fd_ids = ast_fun_ids
        self.is_inner = is_inner
        if not self.is_inner:
            self.check_for_inner_defs()

    def get_inner_func(self, func_name):
        return next(f for f in self.inner_funcs if f.name == func_name)

    @property
    def has_inner_func(self):
        return self.inner_func_idx != []

    @property
    def inner_func_idx(self):
        return self._inner_func_idx

    @inner_func_idx.setter
    def inner_func_idx(self, idx):
        self._inner_func_idx = idx

    @property
    def line_range(self):
        return (self.lineno, self.end_lineno)

    @property
    def inner_funcs(self):
        return self._inner_funcs

    @inner_funcs.setter
    def inner_funcs(self, funcdefs):
        self._inner_funcs = funcdefs

    def get_ho_cls(self, cls_name):
        return next(c for c in self.ho_classes if c.name == cls_name)

    @property
    def has_ho_cls(self):
        return self.ho_cls_idx != []

    @property
    def ho_cls_idx(self):
        return self._ho_cls_idx

    @ho_cls_idx.setter
    def ho_cls_idx(self, idx):
        self._ho_cls_idx = idx

    def set_ho_classes(self):
        if not self.has_ho_cls:
            self.ho_classes = []
        self.ho_classes = [
            HOClsDef(cd=self.body[i], parent_cd=self) for i in self.ho_cls_idx
        ]

    @property
    def ho_classes(self):
        return self._ho_classes

    @ho_classes.setter
    def ho_classes(self, clsdefs):
        self._ho_classes = clsdefs

    # Read-only aliases used in `RecursiveIdSetterMixin.recursively_set_inner_defs`
    intra_func_idx_attr = "_inner_func_idx"
    intra_cls_idx_attr = "_ho_cls_idx"
    has_intra_func = has_inner_func
    has_intra_cls = has_ho_cls
    intra_funcs = inner_funcs
    intra_classes = ho_classes
    intra_func_type_name = "InnerFunc"
    intra_cls_type_name = "HigherOrderClass"

    @property
    def path(self):
        return self.name

class MethodDef(FuncDef):
    """
    Wrap `FuncDef` (in turn wrapping `ast.FunctionDef`) to store a reference to the
    parent classdef's line range on a method.
    """

    def __init__(self, fd, parent_cd):
        self.classes_only = parent_cd.classes_only
        self.parent_name = parent_cd.name
        self.parent_path = parent_cd.path
        self.parent_line_range = parent_cd.line_range
        # moved to end to match ClsDef subclasses
        super().__init__(fd, self.classes_only, is_inner=True)

    @property
    def path(self):
        parent_path = self.parent_path
        return f"{parent_path}.{self.name}"

    @property
    def all_ns_cd_ids(self):
        return self._all_ns_cd_ids

    @all_ns_cd_ids.setter
    def all_ns_cd_ids(self, id_list):
        self._all_ns_cd_ids = id_list

    @property
    def all_ns_fd_ids(self):
        return self._all_ns_fd_ids

    @all_ns_fd_ids.setter
    def all_ns_fd_ids(self, id_list):
        self._all_ns_fd_ids = id_list

class InnerFuncDef(FuncDef):
    """
    Wrap `FuncDef` (in turn wrapping `ast.FunctionDef`) to store a reference to the
    parent funcdef's line range on an inner function.
    """

    # TODO: Note you could use a mixin to remove the duplicate code between this and all other
    # subclasses of clsdef and funcdef
    def __init__(self, fd, parent_fd):
        self.classes_only = parent_fd.classes_only
        self.parent_name = parent_fd.name
        self.parent_path = parent_fd.path
        self.parent_line_range = parent_fd.line_range
        super().__init__(
            fd, self.classes_only, is_inner=True
        )  # moved to end to match ClsDef subclasses

    @property
    def path(self):
        parent_path = self.parent_path
        return f"{parent_path}:{self.name}"

    @property
    def all_ns_cd_ids(self):
        return self._all_ns_cd_ids

    @all_ns_cd_ids.setter
    def all_ns_cd_ids(self, id_list):
        self._all_ns_cd_ids = id_list

    @property
    def all_ns_fd_ids(self):
        return self._all_ns_fd_ids

    @all_ns_fd_ids.setter
    def all_ns_fd_ids(self, id_list):
        self._all_ns_fd_ids = id_list

class DefTypeEnum(Enum):
    Class = ClsDef
    Func = FuncDef

class IntraDefTypeEnum(Enum):
    Method = MethodDef
    InnerClass = InnerClsDef
    InnerFunc = InnerFuncDef
    HigherOrderClass = HOClsDef

class DefTypeToParentTypeEnum(Enum):
    Method = "Class"
    InnerClass = "Class"
    InnerFunc = "Func"
    HigherOrderClass = "Func"

class DefPathTypeEnum(Enum):
    Class = ClassPath
    Func = FuncPath

class IntraDefPathTypeEnum(Enum):
    Method = MethodPath
    InnerClass = InnerClassPath
    InnerFunc = InnerFuncPath
    HigherOrderClass = HigherOrderClassPath

def parse_mv_funcs(linkfile, trunk):
    """
    mvdefs:  the list of functions to move (string list of function names)
    trunk:   AST body for the file (via `ast.parse(fc).body`)
    report:  whether to print [minimal, readable] 'reporting' output

    Produce a dictionary, `mvdef_names`, whose keys are the list of functions
    to move (i.e. the list `mvdefs` becomes the list of keys of `mvdef_names`),
    and the value of which at each key (for a key `m` which indicates the name
    of one of the functions given in `mvdefs` to move) is another dictionary,
    keyed by the full set of names used in that function (`m`) which rely upon
    import statements (i.e. are not builtin names nor passed as parameters to
    the function, nor assigned in the body of the function), and the value
    of which at each name is a final nested dictionary whose keys are always:
      n:      The [0-based] index of the name's source import statement in the
              AST list of all ast.Import and ast.ImportFrom.
      n_i:    The [0-based] index of the name's source import statement within
              the one or more names imported by the source import statement at
              index n (e.g. for `pi` in `from numpy import e, pi`, `n_i` = 1).
      line:   The [1-based] line number of the import statement as given in its
              corresponding AST entry, which indicates the line number of the
              `import` call, not [necessarily] that of the name (i.e. the name
              may not be located there in the file for multi-line imports).
      import: The path of the import statement, which may contain multiple
              parts conjoined by `.` (e.g. `matplotlib.pyplot`)

    I.e. the dictionary with entries accessed as `mvdef_names.get(m).get(k)`
    for `m` in `mvdefs` and `k` in the subset of AST-identified imported names
    in the function with  if f.name not in mvdefs name `m` in the list of
    function definitions `defs`. This access is handed off to the helper
    function `get_def_names`.

    For the names that were imported but not used, the dictionary is not keyed
    by function (as there are no associated functions), and instead the entries
    are accessed as `nondef_names.get(k)` for `k` in `unused_names`. This access
    is handed off to the helper function `set_nondef_names`.
    """
    mvdefs = linkfile.mvdefs
    get_cls = linkfile.classes_only
    if mvdefs is None:
        mvdefs = []  # prepare for `get_def_names`, don't pass a `None` directly
    report_VERBOSE = False  # Silencing debug print statements
    import_types = [ast.Import, ast.ImportFrom]
    imports = [n for n in trunk if type(n) in import_types]
    ast_funcdefs = [n for n in trunk if type(n) is ast.FunctionDef]
    ast_fd_ids = [f.name for f in ast_funcdefs]
    ast_clsdefs = [n for n in trunk if type(n) is ast.ClassDef]
    ast_defs = ast_clsdefs if get_cls else ast_funcdefs
    ast_cd_ids = [c.name for c in ast_clsdefs]
    def_params = {
        "ast_cls_ids": ast_cd_ids,
        "ast_fun_ids": ast_fd_ids,
        "classes_only": linkfile.classes_only,
    }
    linkfile.ast_funcs = [FuncDef(n, **def_params) for n in ast_funcdefs]
    linkfile.ast_classes = [ClsDef(n, **def_params) for n in ast_clsdefs]
    # Any nodes in the AST that aren't imports or ast_defs are 'extra' (as in 'outside')
    ast_sel_type = ClassDef if get_cls else FunctionDef
    extra = [n for n in trunk if type(n) not in [*import_types, ast_sel_type]]
    # Omit names used outside of function definitions so as not to remove them
    linkfile.set_extradef_names(extra)  # sets extradef_names by walking the extra nodes
    if report_VERBOSE:
        print("extra:", extra, file=stderr)
    import_annos = annotate_imports(imports, report=linkfile.report)
    linkfile.mvdef_names = linkfile.get_def_names(mvdefs, import_annos)
    if report_VERBOSE:
        print("mvdef names:", file=stderr)
        pprint_def_names(linkfile.mvdef_names)
    # ------------------------------------------------------------------------ #
    # Next obtain nonmvdef_names
    linkfile_defs = linkfile.ast_classes if get_cls else linkfile.ast_funcs
    nomvdefs = [f.name for f in linkfile_defs if f.name not in mvdefs]
    linkfile.nonmvdef_names = linkfile.get_def_names(nomvdefs, import_annos)
    if report_VERBOSE:
        print("non-mvdef names:", file=stderr)
        pprint_def_names(nonmvdef_names)
    # ------------------------------------------------------------------------ #
    # Next obtain unused_names
    mv_set = set().union(
        *[linkfile.mvdef_names.get(x).keys() for x in linkfile.mvdef_names]
    )
    nomv_set = set().union(
        *[linkfile.nonmvdef_names.get(x).keys() for x in linkfile.nonmvdef_names]
    )
    unused_names = list(set(list(import_annos[0].keys())) - mv_set - nomv_set)
    linkfile.set_nondef_names(unused_names, import_annos)  # sets nondef_names
    linkfile.set_undef_names()  # set undef_names using nondef_names
    if report_VERBOSE:
        print("non-def names (imported but not used in any function def):")
        pprint_def_names(nondefs, no_funcdef_list=True)
    return

max_stars_repo_name: lmmx/mvdef

max_issues_repo_name: lmmx/mvdef

max_forks_repo_name: lmmx/mvdef

reason: from numpy


text:
from unittest import TestCase

import numpy as np

import cate.util.im.utils as utils

class ArrayTest(TestCase):
    def test_numpy_reduce(self):
        nan = np.nan
        a = np.zeros((8, 6))
        a[0::2, 0::2] = 1.1
        a[0::2, 1::2] = 2.2
        a[1::2, 0::2] = 3.3
        a[1::2, 1::2] = 4.4
        a[6, 4] = np.nan

        self.assertEqual(a.shape, (8, 6))
        np.testing.assert_equal(a, np.array([[1.1, 2.2, 1.1, 2.2, 1.1, 2.2],
                                             [3.3, 4.4, 3.3, 4.4, 3.3, 4.4],
                                             [1.1, 2.2, 1.1, 2.2, 1.1, 2.2],
                                             [3.3, 4.4, 3.3, 4.4, 3.3, 4.4],
                                             [1.1, 2.2, 1.1, 2.2, 1.1, 2.2],
                                             [3.3, 4.4, 3.3, 4.4, 3.3, 4.4],
                                             [1.1, 2.2, 1.1, 2.2, nan, 2.2],
                                             [3.3, 4.4, 3.3, 4.4, 3.3, 4.4]]))

        b = utils.downsample_ndarray(a)
        self.assertEqual(b.shape, (4, 3))
        np.testing.assert_equal(b, np.array([[2.75, 2.75, 2.75],
                                             [2.75, 2.75, 2.75],
                                             [2.75, 2.75, 2.75],
                                             [2.75, 2.75, nan]]))

        b = utils.downsample_ndarray(a, aggregator=utils.aggregate_ndarray_first)
        self.assertEqual(b.shape, (4, 3))
        np.testing.assert_equal(b, np.array([[1.1, 1.1, 1.1],
                                             [1.1, 1.1, 1.1],
                                             [1.1, 1.1, 1.1],
                                             [1.1, 1.1, nan]]))

class GetChunkSizeTest(TestCase):
    def test_any_obj(self):
        any_obj = object()
        self.assertEqual(utils.get_chunk_size(any_obj), None)

    def test_xarray_var(self):
        class X:
            pass

        xarray_var = X()
        xarray_var.encoding = dict(chunksizes=(1, 256, 256))
        self.assertEqual(utils.get_chunk_size(xarray_var), (1, 256, 256))

    def test_netcdf4_var(self):
        class X:
            pass

        netcdf4_var = X()
        netcdf4_var.chunks = (1, 900, 1800)
        self.assertEqual(utils.get_chunk_size(netcdf4_var), (1, 900, 1800))

ext: py

max_stars_repo_name: strawpants/cate

max_issues_repo_name: strawpants/cate

max_forks_repo_name: strawpants/cate

reason: import numpy


text:
from typing import Iterator, List, Any

import numpy as np

from . import config_fields
from . import model_helper
from .data_builder import XGBoostData
from .env_parser import extract_dist_env
from .rabit_context import RabitContext
from .data_units import XGBoostResult

class XGBoostPredictor:
    def __init__(self, conf: config_fields.PredictFields):
        col_conf = conf.data_conf.column_format
        self._ret_cols = col_conf.result_columns
        assert self._ret_cols, 'info of ResultColumns is not set!'
        assert self._ret_cols.result_column, 'name of result_column should not be None!'
        assert self._ret_cols.result_column not in col_conf.append_columns, \
            'Name of result_column is duplicated with one of the append columns!'

        self._predict_leaf = self._ret_cols.leaf_column is not None

        model = model_helper.load_launcher_model(conf.model_conf)
        self._booster = model.booster

        if model.meta.params.objective == 'binary:logistic':
            self._converter = binary_logistic_converter
        elif model.meta.params.objective == 'multi:softprob':
            self._converter = multi_softprob_converter
        else:
            self._converter = common_converter

        self._rabit_ctx = RabitContext(extract_dist_env())

    def __enter__(self):
        self._rabit_ctx.__enter__()
        self._rabit_ctx.client_init()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._rabit_ctx.__exit__(exc_type, exc_val, exc_tb)

    def predict(self, data: XGBoostData) -> Iterator[XGBoostResult]:
        batch_ret = self._booster.predict(data.d_matrix)
        batch_ret = self._converter(batch_ret)
        if data.append_info:
            batch_ret.set_append_info(data.append_info)
        if self._predict_leaf:
            batch_ret.set_leaf(self._booster.predict(data.d_matrix, pred_leaf=True))
        return batch_ret.to_iterator()

class XGBoostResultBatch:
    def __init__(self, result=None, prob=None, detail=None):
        assert isinstance(result, np.ndarray)
        assert len(result.shape) == 1, 'result should be 1-d array!'
        self.num_rows = result.shape[0]
        self.result = result
        if prob is not None:
            assert isinstance(prob, np.ndarray)
            assert len(prob.shape) == 1, 'prob should be 1-d array!'
            assert prob.shape[0] == self.num_rows, 'prob should has %d rows!' % self.num_rows
        self.prob = prob
        if detail is not None:
            assert isinstance(detail, np.ndarray)
            assert len(detail.shape) == 2, 'detail should be 2-d array'
            assert detail.shape[0] == self.num_rows, 'detail should has %d rows!' % self.num_rows
        self.detail = detail
        self.leaf = None
        self.append_info = None

    def set_leaf(self, leaf_result: np.ndarray):
        assert isinstance(leaf_result, np.ndarray)
        assert len(leaf_result.shape) == 2
        assert leaf_result.shape[0] == self.num_rows
        self.leaf = leaf_result

    def set_append_info(self, append_info: List[List[Any]]):
        assert isinstance(append_info[0], list)
        assert len(append_info) == self.num_rows
        self.append_info = append_info

    def to_iterator(self) -> Iterator[XGBoostResult]:
        for i in range(self.num_rows):
            yield XGBoostResult(
                result=self.result[i],
                classification_prob=self.prob[i] if self.prob is not None else None,
                classification_detail=self.detail[i].tolist() if self.detail is not None else None,
                leaf_indices=self.leaf[i].tolist() if self.leaf is not None else None,
                append_info=self.append_info[i] if self.append_info is not None else None)

def common_converter(raw_result: np.ndarray):
    return XGBoostResultBatch(raw_result, None, None)

def binary_logistic_converter(raw_result: np.ndarray):
    detail = np.column_stack([1.0 - raw_result, raw_result])
    result = detail.argmax(axis=1)
    prob = detail.max(axis=1)
    return XGBoostResultBatch(result, prob, detail)

def multi_softprob_converter(raw_result: np.ndarray):
    detail = raw_result
    result = detail.argmax(axis=1)
    prob = raw_result.max(axis=1)
    return XGBoostResultBatch(result, prob, detail)

max_stars_repo_name: alipay/ant-xgboost

max_issues_repo_name: alipay/ant-xgboost

max_forks_repo_name: alipay/ant-xgboost


text:
# coding: utf-8
from __future__ import print_function
import numpy as np
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import tensorflow as tf

def main(_):
    csv_file_name = './data/period_trend.csv'
    reader = tf.contrib.timeseries.CSVReader(csv_file_name)
    train_input_fn = tf.contrib.timeseries.RandomWindowInputFn(reader, batch_size=16, window_size=16)
    with tf.Session() as sess:
        data = reader.read_full()
        coord = tf.train.Coordinator()
        tf.train.start_queue_runners(sess=sess, coord=coord)
        data = sess.run(data)
        coord.request_stop()

    ar = tf.contrib.timeseries.ARRegressor(
        periodicities=100, input_window_size=10, output_window_size=6,
        num_features=1,
        loss=tf.contrib.timeseries.ARModel.NORMAL_LIKELIHOOD_LOSS)

    ar.train(input_fn=train_input_fn, steps=1000)

    evaluation_input_fn = tf.contrib.timeseries.WholeDatasetInputFn(reader)
    # keys of evaluation: ['covariance', 'loss', 'mean', 'observed', 'start_tuple', 'times', 'global_step']
    evaluation = ar.evaluate(input_fn=evaluation_input_fn, steps=1)

    (predictions,) = tuple(ar.predict(
        input_fn=tf.contrib.timeseries.predict_continuation_input_fn(
            evaluation, steps=250)))

    plt.figure(figsize=(15, 5))
    plt.plot(data['times'].reshape(-1), data['values'].reshape(-1), label='origin')
    plt.plot(evaluation['times'].reshape(-1), evaluation['mean'].reshape(-1), label='evaluation')
    plt.plot(predictions['times'].reshape(-1), predictions['mean'].reshape(-1), label='prediction')
    plt.xlabel('time_step')
    plt.ylabel('values')
    plt.legend(loc=4)
    plt.savefig('predict_result.jpg')

if __name__ == '__main__':
    tf.logging.set_verbosity(tf.logging.INFO)
    tf.app.run()


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:

import numpy as np
import nibabel as nb
from scipy.interpolate import interpn
from datetime import datetime as dt

from niworkflows.nipype import logging
LOGGER = logging.getLogger('interfaces')

class BSplineFieldmap(object):
    """
    A fieldmap representation object using BSpline basis
    """

    def __init__(self, fmapnii, weights=None, knots_zooms=None, padding=3,
                 pe_dir=1, njobs=-1):

        self._pedir = pe_dir
        if knots_zooms is None:
            knots_zooms = [40., 40., 18.]
            knots_zooms[pe_dir] = 60.

        if not isinstance(knots_zooms, (list, tuple)):
            knots_zooms = [knots_zooms] * 3

        self._knots_zooms = np.array(knots_zooms)

        if isinstance(fmapnii, str):
            fmapnii = nb.as_closest_canonical(nb.load(fmapnii))

        self._fmapnii = fmapnii
        self._padding = padding

        # Pad data with zeros
        self._data = np.zeros(tuple(np.array(
            self._fmapnii.get_data().shape) + 2 * padding))

        # The list of ijk coordinates
        self._fmapijk = get_ijk(self._data)

        # Find padding coordinates
        self._data[padding:-padding,
                   padding:-padding,
                   padding:-padding] = 1
        self._frameijk = self._data[tuple(self._fmapijk.T)] > 0

        # Set data
        self._data[padding:-padding,
                   padding:-padding,
                   padding:-padding] = fmapnii.get_data()

        # Get ijk in homogeneous coords
        ijk_h = np.hstack((self._fmapijk, np.array([1.0] * len(self._fmapijk))[..., np.newaxis]))

        # The list of xyz coordinates
        self._fmapaff = compute_affine(self._data, self._fmapnii.header.get_zooms())
        self._fmapxyz = self._fmapaff.dot(ijk_h.T)[:3, :].T

        # Mask coordinates
        self._weights = self._set_weights(weights)

        # Generate control points
        self._generate_knots()

        self._X = None
        self._coeff = None
        self._smoothed = None

        self._Xinv = None
        self._inverted = None
        self._invcoeff = None

        self._njobs = njobs

    def _generate_knots(self):
        extent = self._fmapaff[:3, :3].dot(self._data.shape[:3])
        self._knots_shape = (np.ceil(
            (extent - self._knots_zooms) / self._knots_zooms) + 3).astype(int)
        self._knots_grid = np.zeros(tuple(self._knots_shape), dtype=np.float32)
        self._knots_aff = compute_affine(self._knots_grid, self._knots_zooms)

        self._knots_ijk = get_ijk(self._knots_grid)
        knots_ijk_h = np.hstack((self._knots_ijk, np.array(
            [1.0] * len(self._knots_ijk))[..., np.newaxis]))  # In homogeneous coords

        # The list of xyz coordinates
        self._knots_xyz = self._knots_aff.dot(knots_ijk_h.T)[:3, :].T

    def _set_weights(self, weights=None):
        if weights is not None:
            extweights = np.ones_like(self._data)
            extweights[self._padding:-self._padding,
                       self._padding:-self._padding,
                       self._padding:-self._padding] = weights

            return extweights[tuple(self._fmapijk.T)]
        return np.ones((len(self._fmapxyz)))

    def _evaluate_bspline(self):
        """ Calculates the design matrix """
        print('[%s] Evaluating tensor-product cubic BSpline on %d points, %d control points' %
              (dt.now(), len(self._fmapxyz), len(self._knots_xyz)))
        self._X = tbspl_eval(self._fmapxyz, self._knots_xyz, self._knots_zooms, njobs=self._njobs)
        print('[%s] Finished BSpline evaluation' % dt.now())

    def fit(self):
        self._evaluate_bspline()

        fieldata = self._data[tuple(self._fmapijk.T)]

        print('[%s] Starting least-squares fitting using %d unmasked points' %
              (dt.now(), len(fieldata[self._weights > 0.0])))
        self._coeff = np.linalg.lstsq(
            self._X[self._weights > 0.0, ...].toarray(),
            fieldata[self._weights > 0.0])[0]
        print('[%s] Finished least-squares fitting' % dt.now())

    def get_coeffmap(self):
        self._knots_grid[tuple(self._knots_ijk.T)] = self._coeff
        return nb.Nifti1Image(self._knots_grid, self._knots_aff, None)

    def get_smoothed(self):
        self._smoothed = np.zeros_like(self._data)
        coords = tuple(self._fmapijk[self._frameijk].T)
        self._smoothed[coords] = self._X[self._frameijk].dot(self._coeff)

        output_image = self._smoothed[self._padding:-self._padding,
                                      self._padding:-self._padding,
                                      self._padding:-self._padding]
        return nb.Nifti1Image(output_image, self._fmapnii.affine, self._fmapnii.header)

    def invert(self):
        targets = self._fmapxyz.copy()
        targets[:, self._pedir] += self._smoothed[tuple(self._fmapijk.T)]
        print('[%s] Inverting transform :: evaluating tensor-product cubic BSpline on %d points, '
              '%d control points' % (dt.now(), len(targets), len(self._knots_xyz)))
        self._Xinv = tbspl_eval(targets, self._knots_xyz, self._knots_zooms, self._njobs)
        print('[%s] Finished BSpline evaluation, %s' %
              (dt.now(), str(self._Xinv.shape)))

        print('[%s] Starting least-squares fitting using %d unmasked points' %
              (dt.now(), len(targets)))
        self._invcoeff = np.linalg.lstsq(
            self._Xinv, self._fmapxyz[:, self._pedir] - targets[:, self._pedir])[0]
        print('[%s] Finished least-squares fitting' % dt.now())

    def get_inverted(self):
        self._inverted = np.zeros_like(self._data)
        self._inverted[tuple(self._fmapijk.T)] = self._X.dot(self._invcoeff)
        return nb.Nifti1Image(self._inverted, self._fmapnii.affine, self._fmapnii.header)

    def interp(self, in_data, inverse=False, fwd_pe=True):
        dshape = tuple(in_data.shape)
        gridxyz = self._fmapxyz.reshape((dshape[0], dshape[1], dshape[2], -1))

        x = gridxyz[:, 0, 0, 0]
        y = gridxyz[0, :, 0, 1]
        z = gridxyz[0, 0, :, 2]

        targets = self._fmapxyz.copy()

        if inverse:
            factor = 1.0 if fwd_pe else -1.0
            targets[:, self._pedir] += factor * \
                self._inverted[tuple(self._fmapijk.T)]
        else:
            targets[:, self._pedir] += self._smoothed[tuple(self._fmapijk.T)]

        interpolated = np.zeros_like(self._data)
        interpolated[tuple(self._fmapijk.T)] = interpn(
            (x, y, z), in_data, [tuple(v) for v in targets],
            bounds_error=False, fill_value=0)

        return nb.Nifti1Image(interpolated, self._fmapnii.affine, self._fmapnii.header)

#     def xfm_coords(self, in_coord):
#         X = fif.tbspl_eval(np.array([in_coord]), self._knots_xyz, self._knots_zooms)
#         new_coord = in_coord + X.dot(self._coeff if not inverse else self._invcoeff)

def get_ijk(data, offset=0):
    """
    Calculates voxel coordinates from data
    """
    from numpy import mgrid

    if not isinstance(offset, (list, tuple)):
        offset = [offset] * 3

    grid = mgrid[offset[0]:(offset[0] + data.shape[0]),
                 offset[1]:(offset[1] + data.shape[1]),
                 offset[2]:(offset[2] + data.shape[2])]
    return grid.reshape(3, -1).T

def compute_affine(data, zooms):
    """
    Compose a RAS affine mat, since the affine of the image might not be RAS
    """
    aff = np.eye(4) * (list(zooms) + [1])
    aff[:3, 3] -= aff[:3, :3].dot(np.array(data.shape[:3], dtype=float) - 1.0) * 0.5
    return aff

def _approx(fmapnii, s=14.):
    """
    Slice-wise approximation of a smooth 2D bspline
    credits: http://scipython.com/book/chapter-8-scipy/examples/two-dimensional-interpolation-\
    with-scipyinterpolaterectbivariatespline/

    """
    from scipy.interpolate import RectBivariateSpline

    if isinstance(fmapnii, str):
        fmapnii = nb.load(fmapnii)

    if not isinstance(s, (tuple, list)):
        s = np.array([s] * 2)

    data = fmapnii.get_data()
    zooms = fmapnii.header.get_zooms()

    knot_decimate = np.floor(s / np.array(zooms)[:2]).astype(np.uint8)
    knot_space = np.array(zooms)[:2] * knot_decimate

    xmax = 0.5 * data.shape[0] * zooms[0]
    ymax = 0.5 * data.shape[1] * zooms[1]

    x = np.arange(-xmax, xmax, knot_space[0])
    y = np.arange(-ymax, ymax, knot_space[1])

    x2 = np.arange(-xmax, xmax, zooms[0])
    y2 = np.arange(-ymax, ymax, zooms[1])

    coeffs = []
    nslices = data.shape[-1]
    for k in range(nslices):
        data2d = data[..., k]
        data2dsubs = data2d[::knot_decimate[0], ::knot_decimate[1]]
        interp_spline = RectBivariateSpline(x, y, data2dsubs)

        data[..., k] = interp_spline(x2, y2)
        coeffs.append(interp_spline.get_coeffs().reshape(data2dsubs.shape))

    # Save smoothed data
    hdr = fmapnii.header.copy()
    caff = fmapnii.affine
    datanii = nb.Nifti1Image(data.astype(np.float32), caff, hdr)

    # Save bspline coeffs
    caff[0, 0] = knot_space[0]
    caff[1, 1] = knot_space[1]
    coeffnii = nb.Nifti1Image(np.stack(coeffs, axis=2), caff, hdr)

    return datanii, coeffnii

def bspl_smoothing(fmapnii, masknii=None, knot_space=[18., 18., 20.]):
    """
    A 3D BSpline smoothing of the fieldmap
    """
    from datetime import datetime as dt

    if not isinstance(knot_space, (list, tuple)):
        knot_space = [knot_space] * 3
    knot_space = np.array(knot_space)

    if isinstance(fmapnii, str):
        fmapnii = nb.load(fmapnii)

    data = fmapnii.get_data()
    zooms = fmapnii.header.get_zooms()

    # Calculate hi-res i
    ijk = np.where(data < np.inf)
    xyz = np.array(ijk).T * np.array(zooms)[np.newaxis, :3]

    # Calculate control points
    xyz_max = xyz.max(axis=0)
    knot_dims = np.ceil(xyz_max / knot_space) + 2
    bspl_grid = np.zeros(tuple(knot_dims.astype(int)))
    bspl_ijk = np.where(bspl_grid == 0)
    bspl_xyz = np.array(bspl_ijk).T * knot_space[np.newaxis, ...]
    bspl_max = bspl_xyz.max(axis=0)
    bspl_xyz -= 0.5 * (bspl_max - xyz_max)[np.newaxis, ...]

    points_ijk = ijk
    points_xyz = xyz

    # Mask if provided
    if masknii is not None:
        if isinstance(masknii, str):
            masknii = nb.load(masknii)
        data[masknii.get_data() <= 0] = 0
        points_ijk = np.where(masknii.get_data() > 0)
        points_xyz = np.array(points_ijk).T * np.array(zooms)[np.newaxis, :3]

    print('[%s] Evaluating tensor-product cubic-bspline on %d points' %
          (dt.now(), len(points_xyz)))
    # Calculate design matrix
    X = tbspl_eval(points_xyz, bspl_xyz, knot_space)
    print('[%s] Finished, bspline grid has %d control points' %
          (dt.now(), len(bspl_xyz)))
    Y = data[points_ijk]

    # Fit coefficients
    print('[%s] Starting least-squares fitting' % dt.now())
    # coeff = (pinv2(X.T.dot(X)).dot(X.T)).dot(Y) # manual way (seems equally
    # slow)
    coeff = np.linalg.lstsq(X, Y)[0]
    print('[%s] Finished least-squares fitting' % dt.now())
    bspl_grid[bspl_ijk] = coeff
    aff = np.eye(4)
    aff[:3, :3] = aff[:3, :3] * knot_space[..., np.newaxis]
    coeffnii = nb.Nifti1Image(bspl_grid, aff, None)

    # Calculate hi-res design matrix:
    # print('[%s] Evaluating tensor-product cubic-bspline on %d points' % (dt.now(), len(xyz)))
    # Xinterp = tbspl_eval(xyz, bspl_xyz, knot_space)
    # print('[%s] Finished, start interpolation' % dt.now())

    # And interpolate
    newdata = np.zeros_like(data)
    newdata[points_ijk] = X.dot(coeff)
    newnii = nb.Nifti1Image(newdata, fmapnii.affine, fmapnii.header)

    return newnii, coeffnii

def tbspl_eval(points, knots, zooms, njobs=None):
    """
    Evaluate tensor product BSpline
    """
    from scipy.sparse import vstack
    from fmriprep.utils.maths import bspl

    points = np.array(points, dtype=float)
    knots = np.array(knots, dtype=float)
    vbspl = np.vectorize(bspl)

    if njobs is not None and njobs < 1:
        njobs = None

    if njobs == 1:
        coeffs = [_evalp((p, knots, vbspl, zooms)) for p in points]
    else:
        from multiprocessing import Pool
        pool = Pool(processes=njobs, maxtasksperchild=100)
        coeffs = pool.map(
            _evalp, [(p, knots, vbspl, zooms) for p in points])
        pool.close()
        pool.join()

    return vstack(coeffs)

def _evalp(args):
    import numpy as np
    from scipy.sparse import csr_matrix

    point, knots, vbspl, zooms = args
    u_vec = (knots - point[np.newaxis, ...]) / zooms[np.newaxis, ...]
    c = vbspl(u_vec.reshape(-1)).reshape((knots.shape[0], 3)).prod(axis=1)
    return csr_matrix(c)

max_stars_repo_name: ericearl/fmriprep

max_issues_repo_name: ericearl/fmriprep

max_forks_repo_name: ericearl/fmriprep

reason: import numpy,from numpy,from scipy


text:
#!/usr/bin/env python

import numpy as np

from ....core.image.affine_image import AffineImage
from ..affine import Affine
from ..histogram_registration import HistogramRegistration
from .._registration import _joint_histogram

from numpy.testing import assert_array_equal
from ....testing import assert_equal, assert_almost_equal, assert_raises

dummy_affine = np.eye(4)

def make_data_uint8(dx=100, dy=100, dz=50):
    return (256 * (np.random.rand(dx, dy, dz)
                   - np.random.rand())).astype('uint8')

def make_data_int16(dx=100, dy=100, dz=50):
    return (256 * (np.random.rand(dx, dy, dz)
                   - np.random.rand())).astype('int16')

def make_data_float64(dx=100, dy=100, dz=50):
    return (256 * (np.random.rand(dx, dy, dz)
                 - np.random.rand())).astype('float64')

def _test_clamping(I, thI=0.0, clI=256):
    R = HistogramRegistration(I, I, from_bins=clI)
    R.subsample(spacing=[1, 1, 1])
    Ic = R._from_data
    Ic2 = R._to_data[1:-1, 1:-1, 1:-1]
    assert_equal(Ic, Ic2)
    dyn = Ic.max() + 1
    assert_equal(dyn, R._joint_hist.shape[0])
    assert_equal(dyn, R._joint_hist.shape[1])
    return Ic, Ic2

def test_clamping_uint8():
    I = AffineImage(make_data_uint8(), dummy_affine, 'ijk')
    _test_clamping(I)

def test_clamping_uint8_nonstd():
    I = AffineImage(make_data_uint8(), dummy_affine, 'ijk')
    _test_clamping(I, 10, 165)

def test_clamping_int16():
    I = AffineImage(make_data_int16(), dummy_affine, 'ijk')
    _test_clamping(I)

def test_clamping_int16_nonstd():
    I = AffineImage(make_data_int16(), dummy_affine, 'ijk')
    _test_clamping(I, 10, 165)

def test_clamping_float64():
    I = AffineImage(make_data_float64(), dummy_affine, 'ijk')
    _test_clamping(I)

def test_clamping_float64_nonstd():
    I = AffineImage(make_data_float64(), dummy_affine, 'ijk')
    _test_clamping(I, 10, 165)

def _test_similarity_measure(simi, val):
    I = AffineImage(make_data_int16(), dummy_affine, 'ijk')
    J = AffineImage(I.get_data().copy(), dummy_affine, 'ijk')
    R = HistogramRegistration(I, J)
    R.subsample(spacing=[2, 1, 3])
    R.similarity = simi
    assert_almost_equal(R.eval(Affine()), val)

def test_correlation_coefficient():
    _test_similarity_measure('cc', 1.0)

def test_correlation_ratio():
    _test_similarity_measure('cr', 1.0)

def test_correlation_ratio_L1():
    _test_similarity_measure('crl1', 1.0)

def test_normalized_mutual_information():
    _test_similarity_measure('nmi', 1.0)

def test_joint_hist_eval():
    I = AffineImage(make_data_int16(), dummy_affine, 'ijk')
    J = AffineImage(I.get_data().copy(), dummy_affine, 'ijk')
    # Obviously the data should be the same
    assert_array_equal(I.get_data(), J.get_data())
    # Instantiate default thing
    R = HistogramRegistration(I, J)
    R.similarity = 'cc'
    null_affine = Affine()
    val = R.eval(null_affine)
    assert_almost_equal(val, 1.0)
    # Try with what should be identity
    R.subsample(spacing=[1, 1, 1])
    assert_array_equal(R._from_data.shape, I.shape)
    val = R.eval(null_affine)
    assert_almost_equal(val, 1.0)

def test_joint_hist_raw():
    # Set up call to joint histogram
    jh_arr = np.zeros((10, 10), dtype=np.double)
    data_shape = (2, 3, 4)
    data = np.random.randint(size=data_shape,
                             low=0, high=10).astype(np.short)
    data2 = np.zeros(np.array(data_shape) + 2, dtype=np.short)
    data2[:] = -1
    data2[1:-1, 1:-1, 1:-1] = data.copy()
    vox_coords = np.indices(data_shape).transpose((1, 2, 3, 0))
    vox_coords = vox_coords.astype(np.double)
    _joint_histogram(jh_arr, data.flat, data2, vox_coords, 0)
    assert_almost_equal(np.diag(np.diag(jh_arr)), jh_arr)

def test_explore():
    I = AffineImage(make_data_int16(), dummy_affine, 'ijk')
    J = AffineImage(make_data_int16(), dummy_affine, 'ijk')
    R = HistogramRegistration(I, J)
    T = Affine()
    simi, params = R.explore(T, (0, [-1, 0, 1]), (1, [-1, 0, 1]))

def test_histogram_registration():
    """ Test the histogram registration class.
    """
    I = AffineImage(make_data_int16(), dummy_affine, 'ijk')
    J = AffineImage(I.get_data().copy(), dummy_affine, 'ijk')
    R = HistogramRegistration(I, J)
    assert_raises(ValueError, R.subsample, spacing=[0, 1, 3])

if __name__ == "__main__":
    import nose
    nose.run(argv=['', __file__])

max_stars_repo_name: neurospin/nipy

max_issues_repo_name: neurospin/nipy

max_forks_repo_name: neurospin/nipy

reason: import numpy,from numpy


text:
import numpy as np
#range
#yas=[1,2,3,5,6,7,7,10,12,13]
yas = [100,13,44,23,56,13,68]
print("Range: ", (np.max(yas)-np.min(yas)))
#variance(kare sapması)
print("Variance: ", (np.var(yas)))
#standar devicion
print("Std: ", (np.std(yas)))

##bir sayının karesini almak
karesi = (2**2)
print("karesi: ",karesi)
#karekök almak
karekok=(np.sqrt(9))
print("kare kökü: ",karekok)

max_stars_repo_name: onselaydin/konumtakip

max_issues_repo_name: onselaydin/konumtakip

max_forks_repo_name: onselaydin/konumtakip


text:
import sys
import time
import os.path
import itertools
import threading
#import cv2
import numpy as np
from skimage import io
#import matplotlib.pyplot as plt

def assignCentroids(x, μ):
	m = x.shape[0]
	c = np.zeros((m,), dtype=np.uint8)
	k = μ.shape[0]
	for i in range(m):
		min_dist = 1000000
		for j in range(k):
			dist = np.linalg.norm(x[i] - μ[j])
			if (dist < min_dist):
				min_dist = dist
				c[i] = j
	return c

def moveCentroids(x, c, k):
	m = x.shape[0]
	n = x.shape[1]
	μ = np.zeros((k, n))
	for i in range(k):
		tmp = np.array([x[ci] for ci in range(m) if c[ci] == i])
		if (tmp.shape[0] == 0):
			continue
		μ[i] = (tmp.sum(axis=0)) / np.float64(tmp.shape[0])
	return μ

def randomInit(x, k):
	#μ = np.array([x[int(np.random.rand() * k)] for i in range(k)])
	#the following is a better way of selecting random datasets from x
	#because using the above statement may select the same data point twice
	tmp = np.copy(x)
	np.random.shuffle(tmp)
	μ = tmp[:k,:]
	return μ

def formClusters(x, k, max_iters):
	#randomly initialze μ
	μ = randomInit(x, k)
	for i in range(max_iters):
		c = assignCentroids(x, μ)
		μ = moveCentroids(x, c, k)
	return μ, c

#here is the animation
def animate():
	for c in itertools.cycle(['|', '/', '-', '\\']):
		if done:
			break
		sys.stdout.write('\rrunning k-means ' + c)
		sys.stdout.flush()
		time.sleep(0.1)
	sys.stdout.write('\rDone!     ')

if __name__ == '__main__':
	try:
		name, n_c = [x for x in input().split()]
		if (os.path.isfile(name)):
			#opencv reads images in bgr format
			rgb = io.imread(name)
			#bgr = cv2.imread(name)
		else:
			raise FileNotFoundError
	except ValueError:
		print("Enter image path and color clusters as 'xyz.png 16'")
		sys.exit(1)
	except FileNotFoundError:
		print("FileNotFoundError : '" + name + "' was not found!")
		sys.exit(1)

	#convert bgr to rgb image
	#rgb = bgr[...,::-1]
	tmp = rgb
	m, n, p = tmp.shape[0], tmp.shape[1], tmp.shape[2]
	x = tmp.reshape((m * n, p))
	done = False
	print('')
	print('This may take a while...Please be patient')
	t = threading.Thread(target=animate)
	t.start()
	μ, c = formClusters(x, int(n_c), 10)
	done = True

	c = c.reshape((m, n))
	#plt.imsave('test_clusters.png', c)
	basename = name[:name.rfind('.')]
	#cv2.imwrite(basename + '_clusters.jpg', c)
	io.imsave(basename + '_clusters.png', c);
	np.save(basename + '_means.npy', μ)
	#plt.imshow(recovered_img)
	#plt.show()


text:
#  Copyright 2021 cms.rendner (Daniel Schmidt)
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
import numpy as np
from pandas import Series, DataFrame

def highlight_even_numbers(data, color="red"):
    attr = f'background-color: {color}'

    if isinstance(data, Series):
        return [attr if v % 2 == 0 else '' for v in data]
    elif isinstance(data, DataFrame):
        return DataFrame(np.where(data % 2 == 0, attr, ''), index=data.index, columns=data.columns)
    else:  # scalar (applymap)
        return attr if data % 2 == 0 else ''

def highlight_max_values(data, chunk_parent=None, **kwargs):
    return highlight_extrema(data, "max", kwargs.get("color", "red"), chunk_parent)

def highlight_min_values(data, chunk_parent=None, **kwargs):
    return highlight_extrema(data, "min", kwargs.get("color", "yellow"), chunk_parent)

def highlight_extrema(data, extrema, color, chunk_parent=None):
    d = data if chunk_parent is None else chunk_parent
    attr = f'background-color: {color}'

    if d.ndim == 1:  # Series from .apply(axis=0) or axis=1
        is_extrema = data == getattr(d, extrema)()
        return [attr if v else '' for v in is_extrema]
    else:  # from .apply(axis=None)
        is_extrema = data == getattr(getattr(d, extrema)(), extrema)()
        return DataFrame(np.where(is_extrema, attr, ''), index=data.index, columns=data.columns)

max_stars_repo_name: rendner/py-plugin-dataframe-viewer

max_issues_repo_name: rendner/py-plugin-dataframe-viewer

max_forks_repo_name: rendner/py-plugin-dataframe-viewer


text:
import pandas as pd
import numpy as np

dataset = pd.read_csv('Foodtruck.csv')

features = dataset.iloc[:, :-1].values
labels = dataset.iloc[:, 1].values

from sklearn.linear_model import LinearRegression
regressor = LinearRegression()
regressor.fit(features, labels)

x=[3.07]
x=np.array(x)
x=x.reshape(1,1)
print(regressor.predict(x))

import pickle

with open('zomato_profit.pkl', 'wb') as file:
    pickle.dump(regressor, file)


text:
import databricks.koalas as ks
import numpy as np
# These function can return a Column Expression or a list of columns expression
# Must return None if the data type can not be handle
import pandas as pd
from pyspark.sql import functions as F

from optimus.engines.base.commons.functions import word_tokenize
from optimus.engines.base.dataframe.functions import DataFrameBaseFunctions
from optimus.engines.base.pandas.functions import PandasBaseFunctions
from optimus.helpers.core import val_to_list, one_list_to_val
from optimus.helpers.raiseit import RaiseIt

class SparkFunctions(PandasBaseFunctions, DataFrameBaseFunctions):
    _engine = ks

    def _to_float(self, series):
        """
        Converts a series values to floats
        """
        return series.astype("float")

    def _to_integer(self, series, default=0):
        """
        Converts a series values to integers
        """
        return series.astype("integer")

    def _to_string(self, series):
        """
        Converts a series values to strings
        """
        return series.astype("str")

    def _to_boolean(self, series):
        """
        Converts a series values to bool
        """
        return series.astype("bool")

    def hist(col_name, df, buckets, min_max=None, dtype=None):
        """
        Create a columns expression to calculate a column histogram
        :param col_name:
        :param df:
        :param buckets:
        :param min_max: Min and max value necessary to calculate the buckets
        :param dtype: Column datatype to calculate the related histogram. Int, String and Dates return different histograms
        :return:
        """
        PYSPARK_NUMERIC_TYPES = ["byte", "short", "big", "int", "double", "float"]
        def is_column_a(df, column, dtypes):
            """
            Check if column match a list of data types
            :param df: dataframe
            :param column: column to be compared with
            :param dtypes: types to be checked
            :return:
            """
            column = val_to_list(column)

            if len(column) > 1:
                RaiseIt.length_error(column, 1)

            data_type = tuple(val_to_list(parse_spark_dtypes(dtypes)))
            column = one_list_to_val(column)

            # Filter columns by data type
            return isinstance(df.schema[column].dataType, data_type)

        def create_exprs(_input_col, _buckets, _func):
            def count_exprs(_exprs):
                return F.sum(F.when(_exprs, 1).otherwise(0))

            _exprs = []
            for i, b in enumerate(_buckets):
                lower = b["lower"]
                upper = b["upper"]

                if is_numeric(lower):
                    lower = round(lower, 2)

                if is_numeric(upper):
                    upper = round(upper, 2)

                if len(_buckets) == 1:
                    count = count_exprs(
                        (_func(_input_col) == lower))
                else:
                    if i == len(_buckets):
                        count = count_exprs(
                            (_func(_input_col) > lower) & (_func(_input_col) <= upper))
                    else:
                        count = count_exprs(
                            (_func(_input_col) >= lower) & (_func(_input_col) < upper))

                info = F.create_map(F.lit("count"), count.cast("int"), F.lit("lower"), F.lit(lower), F.lit("upper"),
                                    F.lit(upper)).alias(
                    "hist_agg" + "_" + _input_col + "_" + str(b["bucket"]))
                _exprs.append(info)
            _exprs = F.array(*_exprs).alias("hist" + _input_col)
            return _exprs

        def hist_numeric(_min_max, _buckets):
            if _min_max is None:
                _min_max = df.agg(F.min(col_name).alias("min"), F.max(col_name).alias("max")).to_dict()[0]

            if _min_max["min"] is not None and _min_max["max"] is not None:
                _buckets = create_buckets(_min_max["min"], _min_max["max"], _buckets)
                _exprs = create_exprs(col_name, _buckets, F.col)
            else:
                _exprs = None

            return _exprs

        def hist_string(_buckets):
            _buckets = create_buckets(0, 50, _buckets)
            func = F.length
            return create_exprs(col_name, _buckets, func)

        def hist_date():
            now = datetime.datetime.now()
            current_year = now.year
            oldest_year = 1950

            # Year
            _buckets = create_buckets(oldest_year, current_year, current_year - oldest_year)
            func = F.year
            year = create_exprs(col_name, _buckets, func)

            # Month
            _buckets = create_buckets(1, 12, 11)
            func = F.month
            month = create_exprs(col_name, _buckets, func)

            # Day
            _buckets = create_buckets(1, 31, 31)
            func = F.dayofweek
            day = create_exprs(col_name, _buckets, func)

            # Hour
            _buckets = create_buckets(0, 23, 23)
            func = F.hour
            hour = create_exprs(col_name, _buckets, func)

            # Min
            _buckets = create_buckets(0, 60, 60)
            func = F.minute
            minutes = create_exprs(col_name, _buckets, func)

            # Second
            _buckets = create_buckets(0, 60, 60)
            func = F.second
            second = create_exprs(col_name, _buckets, func)

            exprs = F.create_map(F.lit("years"), year, F.lit("months"), month, F.lit("weekdays"), day,
                                 F.lit("hours"), hour, F.lit("minutes"), minutes, F.lit("seconds"), second)

            return exprs

        if dtype is not None:
            col_dtype = dtype[col_name]["dtype"]
            if col_dtype == "int" or col_dtype == "decimal":
                exprs = hist_numeric(min_max, buckets)
            elif col_dtype == "string":
                exprs = hist_string(buckets)
            elif col_dtype == "date":
                exprs = hist_date()
            else:
                exprs = None
        else:
            if is_column_a(df, col_name, PYSPARK_NUMERIC_TYPES):
                exprs = hist_numeric(min_max, buckets)

            elif is_column_a(df, col_name, "str"):
                exprs = hist_string(buckets)

            elif is_column_a(df, col_name, "date") or is_column_a(df, col_name, "timestamp"):
                exprs = hist_date()
            else:
                exprs = None

        return exprs

    def create_exprs(_input_col, _buckets, _func):
        def count_exprs(_exprs):
            return F.sum(F.when(_exprs, 1).otherwise(0))

        _exprs = []
        for i, b in enumerate(_buckets):
            lower = b["lower"]
            upper = b["upper"]

            if is_numeric(lower):
                lower = round(lower, 2)

            if is_numeric(upper):
                upper = round(upper, 2)

            if len(_buckets) == 1:
                count = count_exprs(
                    (_func(_input_col) == lower))
            else:
                if i == len(_buckets):
                    count = count_exprs(
                        (_func(_input_col) > lower) & (_func(_input_col) <= upper))
                else:
                    count = count_exprs(
                        (_func(_input_col) >= lower) & (_func(_input_col) < upper))

            info = F.create_map(F.lit("count"), count.cast("int"), F.lit("lower"), F.lit(lower), F.lit("upper"),
                                F.lit(upper)).alias(
                "hist_agg" + "_" + _input_col + "_" + str(b["bucket"]))
            _exprs.append(info)
        _exprs = F.array(*_exprs).alias("hist" + _input_col)
        return _exprs

    @staticmethod
    def dask_to_compatible(dfd):
        from optimus.helpers.converter import dask_dataframe_to_pandas
        return ks.from_pandas(dask_dataframe_to_pandas(dfd))

    @staticmethod
    def new_df(*args, **kwargs):
        return ks.from_pandas(pd.DataFrame(*args, **kwargs))

    @staticmethod
    def df_concat(df_list):
        return ks.concat(df_list, axis=0, ignore_index=True)

    def word_tokenize(self, series):
        return self.to_string(series).map(word_tokenize, na_action=None)

    def count_zeros(self, series, *args):
        return int((self.to_float(series).values == 0).sum())

    def kurtosis(self, series):
        return self.to_float(series).kurtosis()

    def skew(self, series):
        return self.to_float(series).skew()

    def exp(self, series):
        return np.exp(self.to_float(series))

    def sqrt(self, series):
        return np.sqrt(self.to_float(series))

    def reciprocal(self, series):
        return np.reciprocal(self.to_float(series))

    def radians(self, series):
        return np.radians(self.to_float(series))

    def degrees(self, series):
        return np.degrees(self.to_float(series))

    def ln(self, series):
        return np.log(self.to_float(series))

    def log(self, series, base=10):
        return np.log(self.to_float(series)) / np.log(base)

    def sin(self, series):
        return np.sin(self.to_float(series))

    def cos(self, series):
        return np.cos(self.to_float(series))

    def tan(self, series):
        return np.tan(self.to_float(series))

    def asin(self, series):
        return np.arcsin(self.to_float(series))

    def acos(self, series):
        return np.arccos(self.to_float(series))

    def atan(self, series):
        return np.arctan(self.to_float(series))

    def sinh(self, series):
        return np.arcsinh(self.to_float(series))

    def cosh(self, series):
        return np.cosh(self.to_float(series))

    def tanh(self, series):
        return np.tanh(self.to_float(series))

    def asinh(self, series):
        return np.arcsinh(self.to_float(series))

    def acosh(self, series):
        return np.arccosh(self.to_float(series))

    def atanh(self, series):
        return np.arctanh(self.to_float(series))

    def floor(self, series):
        return np.floor(self.to_float(series))

    def ceil(self, series):
        return np.ceil(self.to_float(series))

    def normalize_chars(self, series):
        return series.str.normalize("NFKD").str.encode('ascii', errors='ignore').str.decode('utf8')

    def format_date(self, series, current_format=None, output_format=None):
        return ks.to_datetime(series, format=current_format,
                              errors="coerce").dt.strftime(output_format).reset_index(drop=True)

    def time_between(self, series, value=None, date_format=None):

        value_date_format = date_format

        if is_list_or_tuple(date_format) and len(date_format) == 2:
            date_format, value_date_format = date_format

        if is_list_or_tuple(value) and len(value) == 2:
            value, value_date_format = value

        date = pd.to_datetime(series, format=date_format, errors="coerce")
        value = pd.Timestamp.now() if value is None else pd.to_datetime(value, format=value_date_format,
                                                                        errors="coerce")

        return (value - date)

max_stars_repo_name: ironmussa/Optimus

max_issues_repo_name: ironmussa/Optimus

max_forks_repo_name: ironmussa/Optimus