improve schema

fastparquet/schema.py

@@ -73,10 +73,11 @@ def __init__(self, schema_elements):
             except AttributeError:
                 pass  # already a str
         self.root = schema_elements[0]
+        # repeats loop above doing decoding, could meld
         self.schema_elements_by_name = {
             se[4]: se for se in schema_elements}
         self.tree = {}
-        schema_tree(schema_elements, paths = self.tree)
+        schema_tree(schema_elements, paths=self.tree)
         self._text = None
 
     @property

fastparquet/util.py

@@ -469,6 +469,7 @@ def concat_and_add(*arrs, offset=True):
 
 
 class Task:
+    """Just holds a function and arguments for later call"""
     def __init__(self, func, *args, **kwargs):
         self.func = func
         self.args = args
@@ -479,9 +480,9 @@ def __call__(self):
 
 
 class ThreadPool:
+    """Simplistic fire&forget pool of X threads"""
 
     def __init__(self, num_workers, poll_time=0.0003):
-        import queue
         self.num_workers = num_workers
         self.tasks = []
         self.ntask = 0
@@ -503,6 +504,9 @@ def wait(self):
         while len(self.done) < self.ntask:
             time.sleep(self.poll)
         self.th.clear()
+        self.done.clear()
+        self.tasks.clear()
+        self.ntask = 0
 
     def submit(self, func, *args, **kwargs):
         self.tasks.append(Task(func, *args, **kwargs))
@@ -516,5 +520,5 @@ def go(self, wait=True):
             self.wait()
 
     def run_tasks(self, tasks: list[callable], wait=True):
-        self.tasks = tasks
+        self.tasks.extend(tasks)
         self.go(wait)