Alex C nitpick nitpicked

Author: Alexandre Landeau

python-lib/wordcloud_visualizer.py

@@ -235,6 +235,20 @@ def _normalize_case_token_counts(self, counts: Counter) -> Counter:
         normalized_counts = Counter(dict(zip(df_counts_agg.token_majority_case, df_counts_agg["sum"])))
         return normalized_counts
 
+    def _save_chart(self, fig: plt.figure) -> BytesIO:
+        """Private method to save chart as a bytes stream
+
+        Args:
+            fig (plt.figure): matplotlib figure to save
+
+        Returns:
+            BytesIO: bytes stream containing the chart's data
+        """
+        temp = BytesIO()
+        fig.savefig(temp, bbox_inches=self.bbox_inches, pad_inches=self.pad_inches, dpi=fig.dpi)
+        plt.close()
+        return temp
+
     @time_logging(log_message="Counting tokens")
     def _count_tokens(self, docs: List[Doc]) -> List[Tuple[AnyStr, Dict]]:
         """Private method to count tokens for each document in corpus
@@ -303,19 +317,15 @@ def generate_wordclouds(self, counts: List[Tuple[AnyStr, Dict]]) -> Generator[Tu
                 else:
                     fig = self._generate_wordcloud(frequencies=count, language=self.language, title=wordcloud_title)
                 # Return chart
-                temp = BytesIO()
-                fig.savefig(temp, bbox_inches=self.bbox_inches, pad_inches=self.pad_inches, dpi=fig.dpi)
-                plt.close()
+                temp = self._save_chart(fig)
                 yield (temp, output_file_name)
 
         else:
             # Generate chart
             count = counts[0][1]
             fig = self._generate_wordcloud(frequencies=count, language=self.language)
             # Return chart
-            temp = BytesIO()
-            fig.savefig(temp, bbox_inches=self.bbox_inches, pad_inches=self.pad_inches, dpi=fig.dpi)
-            plt.close()
+            temp = self._save_chart(fig)
             yield (temp, "wordcloud.png")
 
     def tokenize_and_count(self, df: pd.DataFrame) -> List[Tuple[AnyStr, Dict]]: