make the condition more human-readable

Author: jerrypeng7773

src/sagemaker/estimator.py

@@ -696,37 +696,44 @@ def _stage_user_code_in_s3(self) -> str:
         Returns: S3 URI
         """
         if is_pipeline_variable(self.output_path):
-            code_bucket = self.sagemaker_session.default_bucket()
             if self.code_location is None:
+                code_bucket = self.sagemaker_session.default_bucket()
                 code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
-                kms_key = self.output_kms_key
-            else:
-                code_bucket, key_prefix = parse_s3_url(self.code_location)
-                code_s3_prefix = "/".join(
-                    filter(None, [key_prefix, self._current_job_name, "source"])
-                )
-                kms_key = (
-                    self.output_kms_key
-                    if code_bucket == self.sagemaker_session.default_bucket()
-                    else None
-                )
-        else:
-            if self.code_location is None:
-                code_bucket, _ = parse_s3_url(self.output_path)
-                code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
+                kms_key = None
             else:
                 code_bucket, key_prefix = parse_s3_url(self.code_location)
                 code_s3_prefix = "/".join(
                     filter(None, [key_prefix, self._current_job_name, "source"])
                 )
 
+                output_bucket = self.sagemaker_session.default_bucket()
+                kms_key = self.output_kms_key if code_bucket == output_bucket else None
+        else:
             local_mode = self.output_path.startswith("file://")
-
             if local_mode:
-                kms_key = None
+                if self.code_location is None:
+                    code_bucket = self.sagemaker_session.default_bucket()
+                    code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
+                    kms_key = None
+                else:
+                    code_bucket, key_prefix = parse_s3_url(self.code_location)
+                    code_s3_prefix = "/".join(
+                        filter(None, [key_prefix, self._current_job_name, "source"])
+                    )
+                    kms_key = None
             else:
-                output_bucket, _ = parse_s3_url(self.output_path)
-                kms_key = self.output_kms_key if code_bucket == output_bucket else None
+                if self.code_location is None:
+                    code_bucket, _ = parse_s3_url(self.output_path)
+                    code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
+                    kms_key = self.output_kms_key
+                else:
+                    code_bucket, key_prefix = parse_s3_url(self.code_location)
+                    code_s3_prefix = "/".join(
+                        filter(None, [key_prefix, self._current_job_name, "source"])
+                    )
+
+                    output_bucket, _ = parse_s3_url(self.output_path)
+                    kms_key = self.output_kms_key if code_bucket == output_bucket else None
 
         return tar_and_upload_dir(
             session=self.sagemaker_session.boto_session,