Skip to content

Commit

Permalink
fix: Using get_type_hints instead of inspect for udf return type
Browse files Browse the repository at this point in the history
  • Loading branch information
Bhargav Dodla committed Aug 7, 2024
1 parent 125fa49 commit 8838a11
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 3 deletions.
4 changes: 2 additions & 2 deletions sdk/python/feast/on_demand_feature_view.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import inspect
import warnings
from types import FunctionType
from typing import Any, Optional, Union
from typing import Any, Optional, Union, get_type_hints

import dill
import pandas as pd
Expand Down Expand Up @@ -631,7 +631,7 @@ def mainify(obj) -> None:
obj.__module__ = "__main__"

def decorator(user_function):
return_annotation = inspect.signature(user_function).return_annotation
return_annotation = get_type_hints(user_function).get("return", inspect._empty)
udf_string = dill.source.getsource(user_function)
mainify(user_function)
if mode == "pandas":
Expand Down
56 changes: 55 additions & 1 deletion sdk/python/tests/unit/infra/scaffolding/test_repo_operations.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
import os
import tempfile
from contextlib import contextmanager
from pathlib import Path
from tempfile import TemporaryDirectory
Expand All @@ -6,7 +8,13 @@

import assertpy

from feast.repo_operations import get_ignore_files, get_repo_files, read_feastignore
from feast.repo_operations import (
get_ignore_files,
get_repo_files,
parse_repo,
read_feastignore,
)
from tests.utils.cli_repo_creator import CliRunner


@contextmanager
Expand Down Expand Up @@ -140,3 +148,49 @@ def test_feastignore_with_stars2():
(repo_root / "foo1/c.py").resolve(),
]
)


def test_parse_repo():
"Test to ensure that the repo is parsed correctly"
runner = CliRunner()
with tempfile.TemporaryDirectory(dir=os.getcwd()) as temp_dir:
# Make sure the path is absolute by resolving any symlinks
temp_path = Path(temp_dir).resolve()
result = runner.run(["init", "my_project"], cwd=temp_path)
repo_path = Path(temp_path / "my_project" / "feature_repo")
assert result.returncode == 0

repo_contents = parse_repo(repo_path)

assert len(repo_contents.data_sources) == 3
assert len(repo_contents.feature_views) == 2
assert len(repo_contents.on_demand_feature_views) == 2
assert len(repo_contents.stream_feature_views) == 0
assert len(repo_contents.entities) == 2
assert len(repo_contents.feature_services) == 3


def test_parse_repo_with_future_annotations():
"Test to ensure that the repo is parsed correctly when using future annotations"
runner = CliRunner()
with tempfile.TemporaryDirectory(dir=os.getcwd()) as temp_dir:
# Make sure the path is absolute by resolving any symlinks
temp_path = Path(temp_dir).resolve()
result = runner.run(["init", "my_project"], cwd=temp_path)
repo_path = Path(temp_path / "my_project" / "feature_repo")
assert result.returncode == 0

with open(repo_path / "example_repo.py", "r") as f:
existing_content = f.read()

with open(repo_path / "example_repo.py", "w") as f:
f.write("from __future__ import annotations" + "\n" + existing_content)

repo_contents = parse_repo(repo_path)

assert len(repo_contents.data_sources) == 3
assert len(repo_contents.feature_views) == 2
assert len(repo_contents.on_demand_feature_views) == 2
assert len(repo_contents.stream_feature_views) == 0
assert len(repo_contents.entities) == 2
assert len(repo_contents.feature_services) == 3

0 comments on commit 8838a11

Please sign in to comment.