Add assert_type in dataclass transformer

[pingsutw]

Signed-off-by: Kevin Su [email protected]

TL;DR

slack thread

Add assert_type()in dataclass transformer, so we don't need to specify type_hints when using flyte remote.

We got an error when using dataclass as input in the flyte remote without passing type_hints, because guess_python_type is not equal to the user-defined dataclass type

@dataclass_json
@dataclass
class Foo(object):
    a: int = 0

@task
def t1(a: Foo):
    ...

In this example, we use guess_python_type (FooSchema) as our expected_type by default if using flyte remote to execute the above task.
However, FooSchema is created by flytekit and it's not equal to the user-defined dataclass (Foo), so assert_type raises an error.

We can override assert_type in the dataclass transformer, and check that the type of each attribute is the same as the guess_type by iterating all attributes in dataclass.

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

The code to reproduce

from dataclasses import dataclass
from typing import Optional
from dataclasses_json import dataclass_json
from flytekit import task, workflow, dynamic


@dataclass_json
@dataclass
class Args(object):
    max_iter: int = 0
    init_eval: float = 0
    max_eval: str = 0


@dataclass_json
@dataclass
class Schema(object):
    kwargs: Optional[Args] = None


@task
def t11(schema: Schema):
    print(schema)


@workflow
def wf3(schema: Schema = Schema(Args(1, 2.0, "hello"))):
    t11(schema=schema)


if __name__ == "__main__":
    wf3(schema=Schema(Args(1, 2.0, "hello")))

from flytekit.remote import FlyteRemote
from flytekit.configuration import Config
from example_dataclass import Args, Schema

remote = FlyteRemote(
    config=Config.auto("/Users/kevin/.flyte/config-remote.yaml"),
    default_project="flytesnacks",
    default_domain="development",
)

flyte_workflow = remote.fetch_workflow(
    name="example_dataclass.wf3", project="flytesnacks", domain="development"
)

execution = remote.execute(
    flyte_workflow, inputs={"schema": Schema(Args(1, 2, "hello"))}, wait=False,
)

Tracking Issue

flyteorg/flyte#2844

Follow-up issue

NA

[codecov]

Codecov Report

Merging #1149 (3254a53) into master (22580d8) will increase coverage by 0.04%.
The diff coverage is 90.19%.

@@            Coverage Diff             @@
##           master    #1149      +/-   ##
==========================================
+ Coverage   68.38%   68.43%   +0.04%     
==========================================
  Files         288      288              
  Lines       25963    26014      +51     
  Branches     2899     2910      +11     
==========================================
+ Hits        17756    17802      +46     
- Misses       7728     7733       +5     
  Partials      479      479              

Impacted Files	Coverage Δ
flytekit/core/type_engine.py	59.24% <79.16%> (+0.60%)	⬆️
tests/flytekit/unit/core/test_type_engine.py	98.39% <100.00%> (+0.05%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[wild-endeavor]

is there really nothing in dataclass json or json schema that will do this? we have to write our own code?

[wild-endeavor]

is it easy to add a test for this line?

[pingsutw]

Updated it.
I've tried to find a kind of function that can compare the attribute in dataclass, but I didn't find that.