Skip to content

Commit

Permalink
feat(utils/context): rename AssetUniqueKey and AssetAliasUniqueKey to…
Browse files Browse the repository at this point in the history
…_obj methods to to_asset, to_asset_alias
  • Loading branch information
Lee-W committed Dec 9, 2024
1 parent 0a81fa4 commit 7140d60
Show file tree
Hide file tree
Showing 3 changed files with 8 additions and 4 deletions.
4 changes: 3 additions & 1 deletion airflow/models/taskinstance.py
Original file line number Diff line number Diff line change
Expand Up @@ -2765,7 +2765,9 @@ def _register_asset_changes_int(
)
}
if missing_assets := [
unique_key.to_obj() for unique_key, _ in asset_alias_names if unique_key not in asset_models
asset_unique_key.to_asset()
for asset_unique_key, _ in asset_alias_names
if asset_unique_key not in asset_models
]:
asset_models.update(
(AssetUniqueKey.from_asset(asset_obj), asset_obj)
Expand Down
4 changes: 3 additions & 1 deletion airflow/utils/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,7 +208,9 @@ def __str__(self) -> str:
return f"OutletEventAccessors(_dict={self._dict})"

def __iter__(self) -> Iterator[Asset | AssetAlias]:
return iter(key.to_obj() if isinstance(key, BaseAssetUniqueKey) else key for key in self._dict)
return iter(
key.to_asset() if isinstance(key, AssetUniqueKey) else key.to_asset_alias() for key in self._dict
)

def __len__(self) -> int:
return len(self._dict)
Expand Down
4 changes: 2 additions & 2 deletions task_sdk/src/airflow/sdk/definitions/asset/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ class AssetUniqueKey(NamedTuple):
def from_asset(asset: Asset | AssetModel) -> AssetUniqueKey:
return AssetUniqueKey(name=asset.name, uri=asset.uri)

def to_obj(self) -> Asset:
def to_asset(self) -> Asset:
return Asset(name=self.name, uri=self.uri)


Expand All @@ -69,7 +69,7 @@ class AssetAliasUniqueKey(NamedTuple):
def from_asset_alias(asset_alias: AssetAlias) -> AssetAliasUniqueKey:
return AssetAliasUniqueKey(name=asset_alias.name)

def to_obj(self) -> AssetAlias:
def to_asset_alias(self) -> AssetAlias:
return AssetAlias(name=self.name)


Expand Down

0 comments on commit 7140d60

Please sign in to comment.