Files
bearhub/bauh/view/util/disk.py

95 lines
2.8 KiB
Python

import json
import logging
import os
import time
from threading import Thread, Lock
from typing import Type, Dict
import yaml
from bauh.api.abstract.cache import MemoryCache
from bauh.api.abstract.disk import DiskCacheLoader, DiskCacheLoaderFactory
from bauh.api.abstract.model import SoftwarePackage
class AsyncDiskCacheLoader(Thread, DiskCacheLoader):
def __init__(self, cache_map: Dict[Type[SoftwarePackage], MemoryCache], logger: logging.Logger):
super(AsyncDiskCacheLoader, self).__init__(daemon=True)
self.pkgs = []
self._work = True
self.lock = Lock()
self.cache_map = cache_map
self.logger = logger
self.processed = 0
def fill(self, pkg: SoftwarePackage, sync: bool = False):
"""
Adds a package which data must be read from the disk to a queue (if not sync)
:param pkg:
:param sync:
:return:
"""
if pkg and pkg.supports_disk_cache():
if sync:
self._fill_cached_data(pkg)
else:
self.pkgs.append(pkg)
def stop_working(self):
self._work = False
def run(self):
last = 0
while True:
time.sleep(0.00001)
if len(self.pkgs) > self.processed:
pkg = self.pkgs[last]
self._fill_cached_data(pkg)
self.processed += 1
last += 1
elif not self._work:
break
def _fill_cached_data(self, pkg: SoftwarePackage) -> bool:
if os.path.exists(pkg.get_disk_data_path()):
disk_path = pkg.get_disk_data_path()
ext = disk_path.split('.')[-1]
with open(disk_path) as f:
if ext == 'json':
cached_data = json.loads(f.read())
elif ext in {'yml', 'yaml'}:
cached_data = yaml.load(f.read())
else:
raise Exception('The cached data file {} has an unsupported format'.format(disk_path))
if cached_data:
pkg.fill_cached_data(cached_data)
cache = self.cache_map.get(pkg.__class__)
if cache:
cache.add_non_existing(str(pkg.id), cached_data)
return True
return False
class DefaultDiskCacheLoaderFactory(DiskCacheLoaderFactory):
def __init__(self, logger: logging.Logger):
super(DefaultDiskCacheLoaderFactory, self).__init__()
self.logger = logger
self.cache_map = {}
def map(self, pkg_type: Type[SoftwarePackage], cache: MemoryCache):
if pkg_type:
if pkg_type not in self.cache_map:
self.cache_map[pkg_type] = cache
def new(self) -> AsyncDiskCacheLoader:
return AsyncDiskCacheLoader(cache_map=self.cache_map, logger=self.logger)