tests/plugins/test_logging_task_nodes.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

# NOTE: This file contains nodes for the test 'test_logging_parallel_nodes' in
# test_logging_task_nodes.py. They are required to be in a separate file in order to properly
# test the multi-processing executor.

from hamilton.htypes import Collect, Parallelizable
from hamilton.plugins.h_logging import get_logger


def b() -> int:
    return 5


def c(b: int) -> Parallelizable[int]:
    for i in range(b):
        yield i


def d(c: int) -> int:
    logger = get_logger("test_logging_parallel_nodes")
    logger.warning("Context aware message")
    return 2 * c


def e(d: Collect[int]) -> int:
    return sum(d)


def f(e: int) -> int:
    return e