From 5e52db6bca5778523db641594aa150355588bfd6 Mon Sep 17 00:00:00 2001
From: Mark Kuhn <kuhnmar@amazon.com>
Date: Wed, 14 Sep 2022 10:07:58 -0700
Subject: [PATCH] use param in context copy instead of added functn

---
 README.md                                      |  2 +-
 aws_embedded_metrics/logger/metrics_context.py | 15 +++------------
 aws_embedded_metrics/logger/metrics_logger.py  |  3 +--
 3 files changed, 5 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index 699c368..f4b4ac4 100644
--- a/README.md
+++ b/README.md
@@ -167,7 +167,7 @@ set_namespace("MyApplication")
 - **flush**()
 
 Flushes the current MetricsContext to the configured sink and resets all properties and metric values. The namespace and default dimensions will be preserved across flushes.
-Custom dimensions are **not** preserved by default, but this behavior can be changed by invoking `logger.flush_preserve_dimensions = True`, so that custom dimensions would be preserved after each flushing thereafter.
+Custom dimensions are **not** preserved by default, but this behavior can be changed by setting `logger.flush_preserve_dimensions = True`, so that custom dimensions would be preserved after each flushing thereafter.
 
 Example:
 
diff --git a/aws_embedded_metrics/logger/metrics_context.py b/aws_embedded_metrics/logger/metrics_context.py
index 13a30c0..944124a 100644
--- a/aws_embedded_metrics/logger/metrics_context.py
+++ b/aws_embedded_metrics/logger/metrics_context.py
@@ -152,9 +152,10 @@ def get_dimensions(self) -> List[Dict]:
     def __has_default_dimensions(self) -> bool:
         return self.default_dimensions is not None and len(self.default_dimensions) > 0
 
-    def create_copy_with_context(self) -> "MetricsContext":
+    def create_copy_with_context(self, preserve_dimensions: bool = False) -> "MetricsContext":
         """
         Creates a deep copy of the context excluding metrics.
+        Custom dimensions are NOT preserved by default unless preserve_dimensions parameter is set.
         """
         new_properties: Dict = {}
         new_properties.update(self.properties)
@@ -169,7 +170,7 @@ def create_copy_with_context(self) -> "MetricsContext":
         #
         # my_func()
         # my_func()
-        new_dimensions: List[Dict] = []
+        new_dimensions: List[Dict] = [] if not preserve_dimensions else self.dimensions
 
         new_default_dimensions: Dict = {}
         new_default_dimensions.update(self.default_dimensions)
@@ -178,16 +179,6 @@ def create_copy_with_context(self) -> "MetricsContext":
             self.namespace, new_properties, new_dimensions, new_default_dimensions
         )
 
-    def create_copy_with_context_with_dimensions(self) -> "MetricsContext":
-        """
-        Creates a deep copy of the context excluding metrics.
-        Custom dimensions will be copied, this helps with the reuse of dimension sets.
-        """
-        new_context = self.create_copy_with_context()
-        new_context.dimensions.extend(self.dimensions)
-
-        return new_context
-
     @staticmethod
     def empty() -> "MetricsContext":
         return MetricsContext()
diff --git a/aws_embedded_metrics/logger/metrics_logger.py b/aws_embedded_metrics/logger/metrics_logger.py
index e9caa04..0770af9 100644
--- a/aws_embedded_metrics/logger/metrics_logger.py
+++ b/aws_embedded_metrics/logger/metrics_logger.py
@@ -43,8 +43,7 @@ async def flush(self) -> None:
 
         # accept and reset the context
         sink.accept(self.context)
-        self.context = self.context.create_copy_with_context() if not self.flush_preserve_dimensions \
-            else self.context.create_copy_with_context_with_dimensions()
+        self.context = self.context.create_copy_with_context(self.flush_preserve_dimensions)
 
     def __configure_context_for_environment(self, env: Environment) -> None:
         default_dimensions = {