Skip to content

Commit 6e70129

Browse files
authored
Add correct types to Scalar and Primitive constructors in literals.py (#2778)
* Add some type-information to `literals.py` Signed-off-by: Felix Mulder <[email protected]> * Fixup imports Signed-off-by: Felix Mulder <[email protected]> --------- Signed-off-by: Felix Mulder <[email protected]> Signed-off-by: Felix Mulder <[email protected]>
1 parent cc4d27b commit 6e70129

File tree

1 file changed

+22
-21
lines changed

1 file changed

+22
-21
lines changed

Diff for: flytekit/models/literals.py

+22-21
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from datetime import datetime as _datetime
2+
from datetime import timedelta as _timedelta
23
from datetime import timezone as _timezone
34
from typing import Dict, Optional
45

@@ -48,12 +49,12 @@ def from_flyte_idl(cls, pb2_object):
4849
class Primitive(_common.FlyteIdlEntity):
4950
def __init__(
5051
self,
51-
integer=None,
52-
float_value=None,
53-
string_value=None,
54-
boolean=None,
55-
datetime=None,
56-
duration=None,
52+
integer: Optional[int] = None,
53+
float_value: Optional[float] = None,
54+
string_value: Optional[str] = None,
55+
boolean: Optional[bool] = None,
56+
datetime: Optional[_datetime] = None,
57+
duration: Optional[_timedelta] = None,
5758
):
5859
"""
5960
This object proxies the primitives supported by the Flyte IDL system. Only one value can be set.
@@ -77,35 +78,35 @@ def __init__(
7778
self._duration = duration
7879

7980
@property
80-
def integer(self):
81+
def integer(self) -> Optional[int]:
8182
"""
8283
:rtype: int
8384
"""
8485
return self._integer
8586

8687
@property
87-
def float_value(self):
88+
def float_value(self) -> Optional[float]:
8889
"""
8990
:rtype: float
9091
"""
9192
return self._float_value
9293

9394
@property
94-
def string_value(self):
95+
def string_value(self) -> Optional[str]:
9596
"""
9697
:rtype: Text
9798
"""
9899
return self._string_value
99100

100101
@property
101-
def boolean(self):
102+
def boolean(self) -> Optional[bool]:
102103
"""
103104
:rtype: bool
104105
"""
105106
return self._boolean
106107

107108
@property
108-
def datetime(self):
109+
def datetime(self) -> Optional[_datetime]:
109110
"""
110111
:rtype: datetime.datetime
111112
"""
@@ -114,7 +115,7 @@ def datetime(self):
114115
return self._datetime.replace(tzinfo=_timezone.utc)
115116

116117
@property
117-
def duration(self):
118+
def duration(self) -> Optional[_timedelta]:
118119
"""
119120
:rtype: datetime.timedelta
120121
"""
@@ -703,15 +704,15 @@ def from_flyte_idl(cls, pb2_object):
703704
class Scalar(_common.FlyteIdlEntity):
704705
def __init__(
705706
self,
706-
primitive: Primitive = None,
707-
blob: Blob = None,
708-
binary: Binary = None,
709-
schema: Schema = None,
710-
union: Union = None,
711-
none_type: Void = None,
712-
error: Error = None,
713-
generic: Struct = None,
714-
structured_dataset: StructuredDataset = None,
707+
primitive: Optional[Primitive] = None,
708+
blob: Optional[Blob] = None,
709+
binary: Optional[Binary] = None,
710+
schema: Optional[Schema] = None,
711+
union: Optional[Union] = None,
712+
none_type: Optional[Void] = None,
713+
error: Optional[Error] = None,
714+
generic: Optional[Struct] = None,
715+
structured_dataset: Optional[StructuredDataset] = None,
715716
):
716717
"""
717718
Scalar wrapper around Flyte types. Only one can be specified.

0 commit comments

Comments
 (0)