controller_util: Expand typing.
The portage_pb2.Portage.Package message will also need to be able
to be populated. Currently, the two messages have the same fields
so no other changes are needed.
BUG=b:194405049
TEST=None
Change-Id: Icc06a757d0296de24c442b27ce954f62cf48290c
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/3273764
Tested-by: Alex Klein <saklein@chromium.org>
Commit-Queue: Alex Klein <saklein@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
diff --git a/api/controller/controller_util.py b/api/controller/controller_util.py
index 9e769e4..20474b8 100644
--- a/api/controller/controller_util.py
+++ b/api/controller/controller_util.py
@@ -4,6 +4,7 @@
"""Utility functions that are useful for controllers."""
import logging
+from typing import TYPE_CHECKING, Union
from chromite.api.gen.chromite.api import sysroot_pb2
from chromite.api.gen.chromiumos import common_pb2
@@ -14,6 +15,9 @@
from chromite.lib import chroot_lib
from chromite.lib import sysroot_lib
+if TYPE_CHECKING:
+ from chromite.api.gen.chromiumos.build.api import portage_pb2
+
class Error(Exception):
"""Base error class for the module."""
@@ -147,7 +151,8 @@
def serialize_package_info(pkg_info: package_info.PackageInfo,
- pkg_info_msg: common_pb2.PackageInfo):
+ pkg_info_msg: Union[common_pb2.PackageInfo,
+ 'portage_pb2.Portage.Package']):
"""Serialize a PackageInfo object to a PackageInfo proto."""
if not isinstance(pkg_info, package_info.PackageInfo):
# Allows us to swap everything to serialize_package_info, and search the