feat: Allow module_name to be specified with the base logger.

yeesian · copybara-github · commit c01df6681e5f · 2024-03-05T09:05:20.000-08:00
PiperOrigin-RevId: 612860933
diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py
@@ -92,6 +92,8 @@ def log_create_complete(
         cls: Type["VertexAiResourceNoun"],
         resource: proto.Message,
         variable_name: str,
+        *,
+        module_name: str = "aiplatform",
     ):
         """Logs create event is complete.
 
@@ -102,17 +104,23 @@ def log_create_complete(
                 Vertex AI Resource Noun class that is being created.
             resource (proto.Message):
                 Vertex AI Resource proto.Message
-            variable_name (str): Name of variable to use for code snippet
+            variable_name (str):
+                Name of variable to use for code snippet.
+            module_name (str):
+                The module namespace under which the Vertex AI Resource Noun
+                is available. Defaults to `aiplatform`.
         """
         self.info(f"{cls.__name__} created. Resource name: {resource.name}")
         self.info(f"To use this {cls.__name__} in another session:")
-        self.info(f"{variable_name} = aiplatform.{cls.__name__}('{resource.name}')")
+        self.info(f"{variable_name} = {module_name}.{cls.__name__}('{resource.name}')")
 
     def log_create_complete_with_getter(
         self,
         cls: Type["VertexAiResourceNoun"],
         resource: proto.Message,
         variable_name: str,
+        *,
+        module_name: str = "aiplatform",
     ):
         """Logs create event is complete.
 
@@ -123,11 +131,16 @@ def log_create_complete_with_getter(
                 Vertex AI Resource Noun class that is being created.
             resource (proto.Message):
                 Vertex AI Resource proto.Message
-            variable_name (str): Name of variable to use for code snippet
+            variable_name (str):
+                Name of variable to use for code snippet.
+            module_name (str):
+                The module namespace under which the Vertex AI Resource Noun
+                is available. Defaults to `aiplatform`.
         """
         self.info(f"{cls.__name__} created. Resource name: {resource.name}")
         self.info(f"To use this {cls.__name__} in another session:")
-        self.info(f"{variable_name} = aiplatform.{cls.__name__}.get('{resource.name}')")
+        usage_message = f"{module_name}.{cls.__name__}.get('{resource.name}')"
+        self.info(f"{variable_name} = {usage_message}")
 
     def log_delete_with_lro(
         self,
