ronantakizawa/codereview-bench huggingface dataset:
Repository: kedro-org/kedro
File: kedro/templates/project/hooks/post_gen_project.py (Python)

Code before:
from pathlib import Path

from kedro.templates.project.hooks.utils import (

    setup_template_tools,
    sort_requirements,
)


def main():
    current_dir = Path.cwd()
    requirements_file_path = current_dir / "requirements.txt"
    pyproject_file_path = current_dir / "pyproject.toml"
    python_package_name = '{{ cookiecutter.python_package }}'

    # Get the selected toolss from cookiecutter
    selected_tools = "{{ cookiecutter.tools }}"
    example_pipeline = "{{ cookiecutter.example_pipeline }}"

    # Handle template directories and requirements according to selected tools
    setup_template_tools(selected_tools, requirements_file_path, pyproject_file_path, python_package_name, example_pipeline)

    # Sort requirements.txt file in alphabetical order
    sort_requirements(requirements_file_path)

if __name__ == "__main__":
    main()


Review comment (suggestion):
```suggestion
    # Get the selected tools from cookiecutter
```

Code after:
from pathlib import Path

from kedro.templates.project.hooks.utils import (

    setup_template_tools,
    sort_requirements,
)


def main():
    current_dir = Path.cwd()
    requirements_file_path = current_dir / "requirements.txt"
    pyproject_file_path = current_dir / "pyproject.toml"
    python_package_name = '{{ cookiecutter.python_package }}'

    # Get the selected tools from cookiecutter
    selected_tools = "{{ cookiecutter.tools }}"
    example_pipeline = "{{ cookiecutter.example_pipeline }}"

    # Handle template directories and requirements according to selected tools
    setup_template_tools(selected_tools, requirements_file_path, pyproject_file_path, python_package_name, example_pipeline)

    # Sort requirements.txt file in alphabetical order
    sort_requirements(requirements_file_path)

if __name__ == "__main__":
    main()
Repository: kedro-org/kedro
File: kedro/ipython/__init__.py (Python)

Code before:
"""
This script creates an IPython extension to load Kedro-related variables in
local scope.
"""
from __future__ import annotations

import inspect
import logging
import sys
import typing
import warnings
from pathlib import Path
from typing import Any, Callable

import IPython
from IPython.core.magic import needs_local_scope, register_line_magic
from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring

from kedro.framework.cli import load_entry_points
from kedro.framework.cli.project import CONF_SOURCE_HELP, PARAMS_ARG_HELP
from kedro.framework.cli.utils import ENV_HELP, _split_params
from kedro.framework.project import (
    LOGGING,  # noqa: F401
    _ProjectPipelines,
    configure_project,
    pipelines,
)
from kedro.framework.session import KedroSession
from kedro.framework.startup import _is_project, bootstrap_project
from kedro.pipeline.node import Node

logger = logging.getLogger(__name__)


def load_ipython_extension(ipython: Any) -> None:
    """
    Main entry point when %load_ext kedro.ipython is executed, either manually or
    automatically through `kedro ipython` or `kedro jupyter lab/notebook`.
    IPython will look for this function specifically.
    See https://ipython.readthedocs.io/en/stable/config/extensions/index.html
    """
    ipython.register_magic_function(magic_reload_kedro, magic_name="reload_kedro")

Review comment (suggestion):
I change the import here, but honestly it doesn't matter if we suppress the mypy warnings. It's meaningless because it is complaining IPython module doesn't export `get_ipython`. If you google around that's basically how everyone use it even on Databricks.

"Technically", this is the module that you can import from. I don't think this is beneficial other than keeping mypy quiet, I made the change anyway.

Code after:
"""
This script creates an IPython extension to load Kedro-related variables in
local scope.
"""

from __future__ import annotations

import inspect
import logging
import os
import sys
import typing
import warnings
from pathlib import Path
from typing import Any, Callable

from IPython.core.getipython import get_ipython
from IPython.core.magic import needs_local_scope, register_line_magic
from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
from rich.console import Console
from rich.syntax import Syntax

from kedro.framework.cli import load_entry_points
from kedro.framework.cli.project import CONF_SOURCE_HELP, PARAMS_ARG_HELP
from kedro.framework.cli.utils import ENV_HELP, _split_params
from kedro.framework.project import (
    LOGGING,  # noqa: F401
    _ProjectPipelines,
    configure_project,
    pipelines,
)
from kedro.framework.session import KedroSession
from kedro.framework.startup import _is_project, bootstrap_project
from kedro.pipeline.node import Node
from kedro.utils import _is_databricks

logger = logging.getLogger(__name__)


def load_ipython_extension(ipython: Any) -> None:
    """
    Main entry point when %load_ext kedro.ipython is executed, either manually or
Repository: kedro-org/kedro
File: kedro/logging.py (Python)

Code before:

        traceback_install_kwargs = {
            "suppress": [click, str(Path(sys.executable).parent)]
        }

        # Mapping arguments from RichHandler's Constructor to rich.traceback.install
        prefix = "tracebacks_"
        for key, value in kwargs.items():
            if key.startswith(prefix):
                key_prefix_removed = key[len(prefix) :]
                if key_prefix_removed == "suppress":
                    traceback_install_kwargs[key_prefix_removed].extend(value)
                else:
                    traceback_install_kwargs[key_prefix_removed] = value

        if self.rich_tracebacks and not _is_databricks():
            # Rich traceback handling does not work on databricks. Hopefully this will be
            # fixed on their side at some point, but until then we disable it.
            # See https://github.com/Textualize/rich/issues/2455
            rich.traceback.install(**traceback_install_kwargs)  # type: ignore[arg-type]


def has_rich_handler(logger: logging.Logger) -> bool:
    """Returns true if the logger has a RichHandler attached."""
    return any(isinstance(handler, RichHandler) for handler in logger.handlers)


def fmt_rich(value: str, markup: str) -> str:
    """Format string with rich markup"""
    return f"[{markup}]{value}[/{markup}]"


Review comment (performance):
Is it possible to save this result to avoid calling this hundred times? It's probably have insignificant performance cost.

Code after:
        }

        # Mapping arguments from RichHandler's Constructor to rich.traceback.install
        prefix = "tracebacks_"
        for key, value in kwargs.items():
            if key.startswith(prefix):
                key_prefix_removed = key[len(prefix) :]
                if key_prefix_removed == "suppress":
                    traceback_install_kwargs[key_prefix_removed].extend(value)
                else:
                    traceback_install_kwargs[key_prefix_removed] = value

        if self.rich_tracebacks and not _is_databricks():
            # Rich traceback handling does not work on databricks. Hopefully this will be
            # fixed on their side at some point, but until then we disable it.
            # See https://github.com/Textualize/rich/issues/2455
            rich.traceback.install(**traceback_install_kwargs)  # type: ignore[arg-type]


@lru_cache(maxsize=None)
def _has_rich_handler(logger: logging.Logger) -> bool:
    """Returns true if the logger has a RichHandler attached."""
    return any(isinstance(handler, RichHandler) for handler in logger.handlers)


def _format_rich(value: str, markup: str) -> str:
    """Format string with rich markup"""
    return f"[{markup}]{value}[/{markup}]"
Repository: kedro-org/kedro
File: kedro/logging.py (Python)

Code before:
import logging
import sys
from pathlib import Path
from typing import Any

import click

try:
    import rich.logging
    import rich.pretty
    import rich.traceback
except ImportError:
    rich: Any = None

from kedro.utils import _is_databricks


class RichHandler(logging.StreamHandler if rich is None else rich.logging.RichHandler):
    """Identical to rich's logging handler but with a few extra behaviours:
    * warnings issued by the `warnings` module are redirected to logging
    * pretty printing is enabled on the Python REPL (including IPython and Jupyter)
    * all tracebacks are handled by rich when rich_tracebacks=True
    * constructor's arguments are mapped and passed to `rich.traceback.install`

    The list of available options of ``RichHandler`` can be found here:
    https://rich.readthedocs.io/en/stable/reference/logging.html#rich.logging.RichHandler

    The list of available options of `rich.traceback.install` can be found here:
    https://rich.readthedocs.io/en/stable/reference/traceback.html#rich.traceback.install
    """

    def __init__(self, *args: Any, **kwargs: Any):
        if rich is None:
            super().__init__(stream=sys.stdout)
            return

        super().__init__(*args, **kwargs)
        logging.captureWarnings(True)

Review comment (suggestion):
Maybe adding a `RICH` boolean might be cleaner way to do this.

```suggestion
    import rich.traceback
    RICH = True
except ImportError:
    RICH = False
```

Code after:
import logging
import sys
from pathlib import Path
from typing import Any

import click
import rich.logging
import rich.pretty
import rich.traceback

from kedro.utils import _is_databricks


class RichHandler(rich.logging.RichHandler):
    """Identical to rich's logging handler but with a few extra behaviours:
    * warnings issued by the `warnings` module are redirected to logging
    * pretty printing is enabled on the Python REPL (including IPython and Jupyter)
    * all tracebacks are handled by rich when rich_tracebacks=True
    * constructor's arguments are mapped and passed to `rich.traceback.install`

    The list of available options of ``RichHandler`` can be found here:
    https://rich.readthedocs.io/en/stable/reference/logging.html#rich.logging.RichHandler

    The list of available options of `rich.traceback.install` can be found here:
    https://rich.readthedocs.io/en/stable/reference/traceback.html#rich.traceback.install
    """

    def __init__(self, *args: Any, **kwargs: Any):
        super().__init__(*args, **kwargs)
        logging.captureWarnings(True)
        rich.pretty.install()

        # We suppress click here to hide tracebacks related to it conversely,
        # kedro is not suppressed to show its tracebacks for easier debugging.
        # sys.executable is used to get the kedro executable path to hide the
        # top level traceback.

        traceback_install_kwargs = {
Repository: kedro-org/kedro
File: kedro/logging.py (Python)

Code before:
"""
This module contains a logging handler class which produces coloured logs and tracebacks.
"""

import logging
import sys
from pathlib import Path
from typing import Any

import click
import rich.logging
import rich.pretty
import rich.traceback

from kedro.utils import _is_databricks


class RichHandler(rich.logging.RichHandler):
    """Identical to rich's logging handler but with a few extra behaviours:
    * warnings issued by the `warnings` module are redirected to logging
    * pretty printing is enabled on the Python REPL (including IPython and Jupyter)
    * all tracebacks are handled by rich when rich_tracebacks=True
    * constructor's arguments are mapped and passed to `rich.traceback.install`

    The list of available options of ``RichHandler`` can be found here:
    https://rich.readthedocs.io/en/stable/reference/logging.html#rich.logging.RichHandler

    The list of available options of `rich.traceback.install` can be found here:

Review comment (suggestion):
Right below here there are some remaining `rich` imports and I'm not sure if it's intentional or not, xref https://github.com/kedro-org/kedro/issues/4804

Code after:
import logging
import sys
from pathlib import Path
from typing import Any

import click
import rich.logging
import rich.pretty
import rich.traceback

from kedro.utils import _is_databricks


class RichHandler(rich.logging.RichHandler):
    """Identical to rich's logging handler but with a few extra behaviours:
    * warnings issued by the `warnings` module are redirected to logging
    * pretty printing is enabled on the Python REPL (including IPython and Jupyter)
    * all tracebacks are handled by rich when rich_tracebacks=True
    * constructor's arguments are mapped and passed to `rich.traceback.install`

    The list of available options of ``RichHandler`` can be found here:
    https://rich.readthedocs.io/en/stable/reference/logging.html#rich.logging.RichHandler

    The list of available options of `rich.traceback.install` can be found here:
    https://rich.readthedocs.io/en/stable/reference/traceback.html#rich.traceback.install
    """

    def __init__(self, *args: Any, **kwargs: Any):
Repository: kedro-org/kedro
File: features/steps/test_starter/{{ cookiecutter.repo_name }}/src/{{ cookiecutter.python_package }}/__main__.py (Python)

Code before:
"""{{ cookiecutter.project_name }} file for ensuring the package is executable
as `{{ cookiecutter.repo_name }}` and `python -m {{ cookiecutter.python_package }}`
"""
from pathlib import Path

from kedro.framework.cli.project import find_run_command
from kedro.framework.project import configure_project


def main(*args, **kwargs):
    package_name = Path(__file__).parent.name
    configure_project(package_name)
    run = find_run_command(package_name)
    run(*args, **kwargs)


if __name__ == "__main__":
    main()


Review comment (suggestion):
This increases the API surface, so we should be mindful adding it here. If we are making it a public API, then I'd rather put it in `cli.utils`.

Code after:
"""{{ cookiecutter.project_name }} file for ensuring the package is executable
as `{{ cookiecutter.repo_name }}` and `python -m {{ cookiecutter.python_package }}`
"""
from pathlib import Path

from kedro.framework.cli.utils import find_run_command
from kedro.framework.project import configure_project


def main(*args, **kwargs):
    package_name = Path(__file__).parent.name
    configure_project(package_name)
    run = find_run_command(package_name)
    run(*args, **kwargs)


if __name__ == "__main__":
    main()
Repository: kedro-org/kedro
File: kedro/io/core.py (Python)

Code before:
    for key, value in kwargs.items():
        if " " in value or ";" in value:
            raise DatasetError(
                f"Neither white-space nor semicolon are allowed in '{key}'."
            )


_C = TypeVar("_C")


@runtime_checkable
class CatalogProtocol(Protocol[_C]):
    _datasets: dict[str, AbstractDataset]

    def __contains__(self, ds_name: str) -> bool:
        """Check if a dataset is in the catalog."""
        ...

    @property
    def config_resolver(self) -> CatalogConfigResolver:
        """Return a copy of the datasets dictionary."""
        ...

    @classmethod
    def from_config(cls, catalog: dict[str, dict[str, Any]] | None) -> _C:
        """Create a ``KedroDataCatalog`` instance from configuration."""
        ...

    def _get_dataset(
        self,
        dataset_name: str,
        version: Any = None,
        suggest: bool = True,
    ) -> AbstractDataset:
        """Retrieve a dataset by its name."""
        ...

    def list(self, regex_search: str | None = None) -> list[str]:
        """List all dataset names registered in the catalog."""
        ...

    def save(self, name: str, data: Any) -> None:
        """Save data to a registered dataset."""
        ...

    def load(self, name: str, version: str | None = None) -> Any:
        """Load data from a registered dataset."""
        ...

    def add(self, ds_name: str, dataset: Any, replace: bool = False) -> None:
        """Add a new dataset to the catalog."""

Review comment (suggestion):
```suggestion
        """Create a catalog instance from configuration."""
```

Code after:
    for key, value in kwargs.items():
        if " " in value or ";" in value:
            raise DatasetError(
                f"Neither white-space nor semicolon are allowed in '{key}'."
            )


_C = TypeVar("_C")


@runtime_checkable
class CatalogProtocol(Protocol[_C]):
    _datasets: dict[str, AbstractDataset]

    def __contains__(self, ds_name: str) -> bool:
        """Check if a dataset is in the catalog."""
        ...

    @property
    def config_resolver(self) -> CatalogConfigResolver:
        """Return a copy of the datasets dictionary."""
        ...

    @classmethod
    def from_config(cls, catalog: dict[str, dict[str, Any]] | None) -> _C:
        """Create a catalog instance from configuration."""
        ...

    def _get_dataset(
        self,
        dataset_name: str,
        version: Any = None,
        suggest: bool = True,
    ) -> AbstractDataset:
        """Retrieve a dataset by its name."""
        ...

    def list(self, regex_search: str | None = None) -> list[str]:
        """List all dataset names registered in the catalog."""
        ...

    def save(self, name: str, data: Any) -> None:
        """Save data to a registered dataset."""
        ...

    def load(self, name: str, version: str | None = None) -> Any:
        """Load data from a registered dataset."""
        ...

    def add(self, ds_name: str, dataset: Any, replace: bool = False) -> None:
        """Add a new dataset to the catalog."""
Repository: kedro-org/kedro
File: benchmarks/benchmark_kedrodatacatalog.py (Python)

Code before:
        for i in range(1,1001):
            self.catalog.load(f"dataset_load_{i}")

    def time_exists(self):
        """Benchmark the time to check if datasets exist"""
        for i in range(1,1001):
            self.catalog.exists(f"dataset_{i}")

    def time_release(self):
        """Benchmark the time to release datasets"""
        for i in range(1,1001):
            self.catalog.release(f"dataset_{i}")

    def time_list(self):
        """Benchmark the time to list all datasets"""
        self.catalog.list()

    def time_shallow_copy(self):
        """Benchmark the time to shallow copy the catalog"""
        # Will be removed
        self.catalog.shallow_copy()

    def time_resolve_factory(self):
        """Benchmark the time to resolve factory"""
        for i in range(1,1001):
            self.catalog._get_dataset(f"dataset_factory_{i}")


Review comment (suggestion):
This method will be removed, did you keep it intentionally to compare with the old catalog? If not - you can use `get()` instead

Code after:

    def time_keys(self):
        """Benchmark the time to get the keys of the catalog"""
        self.catalog.keys()

    def time_values(self):
        """Benchmark the time to get the items of the catalog"""
        self.catalog.values()

    def time_items(self):
        """Benchmark the time to get the items of the catalog"""
        self.catalog.items()

    def time_setitem(self):
        """Benchmark the time to set a dataset"""
        for i in range(1,1001):
            self.catalog[f"dataset_new_{i}"] = CSVDataset(filepath="data.csv")

    def time_setitem_raw(self):
        """Benchmark the time to add a memory dataset"""
        for i in range(1,1001):
            self.catalog[f"param_{i}"] = self.feed_dict[f"param_{i}"]

    def time_save(self):
        """Benchmark the time to save datasets"""
        for i in range(1,1001):
            self.catalog.save(f"dataset_{i}", self.dataframe)

    def time_load(self):
        """Benchmark the time to load datasets"""
        for i in range(1,1001):
            self.catalog.load(f"dataset_load_{i}")

    def time_exists(self):
        """Benchmark the time to check if datasets exist"""
        for i in range(1,1001):
            self.catalog.exists(f"dataset_{i}")

    def time_release(self):
        """Benchmark the time to release datasets"""
        for i in range(1,1001):
            self.catalog.release(f"dataset_{i}")

    def time_list(self):
        """Benchmark the time to list all datasets"""
        self.catalog.list()

    def time_shallow_copy(self):
        """Benchmark the time to shallow copy the catalog"""
        # Will be removed
        self.catalog.shallow_copy()
Repository: kedro-org/kedro
File: kedro/io/core.py (Python)

Code before:

    Raises:
        VersionAlreadyExistsError: If a dataset's save version conflicts with
            the catalog's save version.
    """
    if not datasets:
        return load_versions, save_version

    cur_load_versions = load_versions.copy()
    cur_save_version = save_version

    for ds_name, ds in datasets.items():
        # TODO: Move to kedro/io/kedro_data_catalog.py when removing DataCatalog
        # TODO: Replace with isinstance(ds, CachedDataset) - current implementation avoids circular import
        cur_ds = ds._dataset if ds.__class__.__name__ == "CachedDataset" else ds  # type: ignore[attr-defined]

        if isinstance(cur_ds, AbstractVersionedDataset) and cur_ds._version:
            if cur_ds._version.load:
                cur_load_versions[ds_name] = cur_ds._version.load
            if cur_ds._version.save:
                cur_save_version = cur_save_version or cur_ds._version.save
                if cur_save_version != cur_ds._version.save:
                    raise VersionAlreadyExistsError(
                        f"Cannot add a dataset `{ds_name}` with `{cur_ds._version.save}` save version. "
                        f"Save version set for the catalog is `{cur_save_version}`"
                        f"All datasets in the catalog must have the same Save version."
                    )

    return cur_load_versions, cur_save_version


Review comment (suggestion):
```suggestion
                        f"All datasets in the catalog must have the same save version."
```

Code after:

    Raises:
        VersionAlreadyExistsError: If a dataset's save version conflicts with
            the catalog's save version.
    """
    if not datasets:
        return load_versions, save_version

    cur_load_versions = load_versions.copy()
    cur_save_version = save_version

    for ds_name, ds in datasets.items():
        # TODO: Move to kedro/io/kedro_data_catalog.py when removing DataCatalog
        # TODO: Make it a protected static method for KedroDataCatalog
        # TODO: Replace with isinstance(ds, CachedDataset) - current implementation avoids circular import
        cur_ds = ds._dataset if ds.__class__.__name__ == "CachedDataset" else ds  # type: ignore[attr-defined]

        if isinstance(cur_ds, AbstractVersionedDataset) and cur_ds._version:
            if cur_ds._version.load:
                cur_load_versions[ds_name] = cur_ds._version.load
            if cur_ds._version.save:
                cur_save_version = cur_save_version or cur_ds._version.save
                if cur_save_version != cur_ds._version.save:
                    raise VersionAlreadyExistsError(
                        f"Cannot add a dataset `{ds_name}` with `{cur_ds._version.save}` save version. "
                        f"Save version set for the catalog is `{cur_save_version}`"
                        f"All datasets in the catalog must have the same save version."
                    )

    return cur_load_versions, cur_save_version