Merge pull request #12 from InnoFang/dev

⚡ 🎨 improve the performance of CountableFileIterator

Author: InnoFang

code_counter/core/countable/iterator.py

@@ -1,46 +1,31 @@
-from collections.abc import Iterator
-import os
 # !/usr/bin/env python3
 # -*- coding: utf-8  -*-
 
+import os
 from code_counter.core.countable.file import CountableFile
 from code_counter.conf.config import Config
-from collections import deque
 
 
-class CountableFileIterator(Iterator):
-    def __init__(self, path):
+class CountableFileIterator:
+    def __init__(self):
         self._ignore = Config().ignore
         self._suffix = Config().suffix
-        self._file_queue = deque()
-        self.__search(path)
 
-    def __search(self, input_path):
+    def iter(self, input_path):
         if os.path.isdir(input_path):
             files = os.listdir(input_path)
             for file in files:
                 file_path = os.path.join(input_path, file)
                 if os.path.isdir(file_path):
-                    if file_path in self._ignore:
+                    if file in self._ignore:
                         continue
-                    self.__search(file_path)
+                    yield from self.iter(file_path)
                 else:
                     suffix = os.path.splitext(file)[1]
                     if len(suffix) == 0 or suffix[1:] not in self._suffix:
                         continue
-                    file_path = os.path.join(input_path, file)
-                    self._file_queue.append(CountableFile(file_path))
+                    yield CountableFile(file_path)
         elif os.path.isfile(input_path):
             suffix = os.path.splitext(input_path)[1]
             if len(suffix) > 0 and suffix[1:] in self._suffix:
-                self._file_queue.append(CountableFile(input_path))
-
-    def __iter__(self):
-        return self
-
-    def __next__(self) -> CountableFile:
-        if self._file_queue:
-            fc = self._file_queue.popleft()
-            return fc
-        else:
-            raise StopIteration
+                yield CountableFile(input_path)

code_counter/core/counter.py

@@ -59,7 +59,7 @@ def search(self):
 
         for path in input_path:
             if os.path.exists(path):
-                for cf in CountableFileIterator(path):
+                for cf in CountableFileIterator().iter(path):
                     cf.count()
                     if self.search_args.verbose:
                         print(cf, file=output_file)
@@ -175,6 +175,6 @@ def visualize(self):
         proptease.set_size('small')
         plt.setp(l_text2, fontproperties=proptease)
         plt.axis('equal')
-        plt.title("Inner Pie: Code Files, Outer Pie: Code Type")
+        plt.title("Inner Pie: Code Files, Outer Pie: Code Lines")
         plt.legend(list(self.result['code'].keys()), title="Abbreviation", loc='best', bbox_to_anchor=(1.05, 1))
         plt.show()