Refactor yaml_list a bit

flexget/components/managed_lists/lists/yaml_list.py

@@ -1,132 +1,54 @@
-from unicodedata import normalize
+import random
 from collections.abc import MutableSet
-import re
-from urllib.parse import quote
+from typing import Optional
 
 from loguru import logger
-from yaml import dump as save_yaml, load as load_yaml
+from yaml import dump as dump_yaml
+from yaml import safe_load as load_yaml
 
 from flexget import plugin
-from flexget.plugin import PluginError
-from flexget.event import event
 from flexget.entry import Entry
-from flexget.utils import json, qualities
-from flexget.task import EntryIterator, EntryContainer
+from flexget.event import event
+from flexget.plugin import PluginError
+from flexget.utils import json
 
 PLUGIN_NAME = 'yaml_list'
 
 logger = logger.bind(name=PLUGIN_NAME)
 
 
-def jsonify(data):
-    """
-    Ensures that data is JSON friendly
-    """
-
-    if isinstance(data, str):
-        return data
-
-    try:
-        _ = (e for e in data)
-    except TypeError:
-        return data
-
-    for item in data:
-        if isinstance(data[item], (EntryIterator, EntryContainer)):
-            lists = list(data[item])
-            new_list = []
-            for lst in lists:
-                dic_list = jsonify(dict(lst))
-                new_list.append(dic_list)
-            data[item] = new_list
-        elif isinstance(data[item], Entry):
-            data[item] = jsonify(dict(data[item]))
-        elif isinstance(data[item], qualities.Quality):
-            data[item] = str(data[item])
-        elif isinstance(data[item], dict):
-            data[item] = jsonify(data[item])
-        else:
-            try:
-                data[item] = json.dumps(data[item])
-                data[item] = json.loads(data[item])
-            except TypeError:
-                del data[item]
-
-    data.pop('_backlog_snapshot', None)
-
-    return data
-
-
-def simplify(text: str) -> str:
-    """ Siplify text """
-
-    if not isinstance(text, str):
-        return text
-
-    # Replace accented chars by their 'normal' couterparts
-    result = normalize('NFKD', text)
-
-    # Symbols that should be converted to white space
-    result = re.sub(r'[ \(\)\-_\[\]\.]+', ' ', result)
-    # Leftovers
-    result = re.sub(r"[^a-zA-Z0-9 ]", "", result)
-    # Replace multiple white spaces with one
-    result = ' '.join(result.split())
-
-    return result
-
-
 class YamlManagedList(MutableSet):
-    _yaml_file = None
-    _yaml_fields = []
-    _yaml_items = {}
+    def __init__(self, path: str, fields: list):
+        self.filename = path
 
-    def __init__(self, path: str, fields, *args, **kwargs):
-        self._yaml_file = path
-
-        if isinstance(fields, list):
-            self._yaml_fields = fields
+        self.fields = fields
 
+        self.entries = []
         try:
-            content = open(self._yaml_file)
-        except FileNotFoundError as e:
-            output_dict = {}
+            content = open(self.filename)
+        except FileNotFoundError as exc:
+            pass
         else:
             try:
-                output_dict = load_yaml(content)
+                # TODO: use the load from our serialization system if that goes in
+                entries = load_yaml(content)
             except Exception as e:
-                raise PluginError(f'Error opening yaml file `{self._yaml_file}`: {e}')
-
-        if isinstance(output_dict, dict):
-            self._yaml_items = output_dict
+                raise PluginError(f'Error opening yaml file `{self.filename}`: {e}')
+        if not entries:
+            return
+        if isinstance(entries, list):
+            for entry in entries:
+                if isinstance(entry, dict):
+                    entry = Entry(**entry)
+                else:
+                    raise PluginError(f'Elements of `{self.filename}` must be dictionaries')
+                if not entry.get('url'):
+                    entry['url'] = f'mock://localhost/entry_list/{random.random()}'
+                self.entries.append(entry)
         else:
-            raise PluginError(f'List `{self._yaml_file}` must be a yaml with objects')
-
-    def _output_item(self, title: str, item: dict) -> dict:
-        """Returns item in output format
-
-        Args:
-            title (str): entry name
-            item (dict): entry data
-
-        Returns:
-            [dict]: Output formated item
-        """
-
-        if not item:
-            return None
-
-        item = self._get_item_fields(item)
-        output_item = {}
-        output_item['title'] = title
-        output_item.update(item)
-
-        if 'url' not in output_item:
-            output_item['url'] = f'mock://{quote(title)}'
+            raise PluginError(f'List `{self.filename}` must be a yaml list')
 
-        return output_item
-
-    def _get_item_fields(self, item: dict) -> dict:
+    def filter_keys(self, item: dict) -> dict:
         """Gets items with limited keys
 
         Args:
@@ -135,28 +57,19 @@ def _get_item_fields(self, item: dict) -> dict:
         Returns:
             dict: Item with limited keys
         """
+        required_fields = ['title']
+        if not self.fields:
+            return {k: item[k] for k in item if not k.startswith('_')}
+        return {k: item[k] for k in item if k in self.fields or k in required_fields}
 
-        output_items = {}
-        for key, data in item.items():
-            if key == 'title':
-                continue
-
-            if self._yaml_fields and key not in self._yaml_fields:
-                continue
-
-            output_items[key] = jsonify(data)
-
-        return output_items
+    def matches(self, entry1, entry2) -> bool:
+        return entry1['title'] == entry2['title']
 
     def __iter__(self):
-        new_list = []
-        for key, data in self._yaml_items.items():
-            new_list.append(self._output_item(key, data))
-
-        return iter(new_list)
+        return iter(self.entries)
 
     def __len__(self):
-        return len(self._yaml_items.keys())
+        return len(self.entries)
 
     def __contains__(self, item):
         return bool(self.get(item))
@@ -168,41 +81,27 @@ def save_yaml(self):
             PluginError: Error
         """
 
+        out = []
+        for entry in self.entries:
+            out.append(json.coerce(self.filter_keys(entry)))
+
         try:
-            with open(self._yaml_file, 'w') as outfile:
-                save_yaml(self._yaml_items, outfile, default_flow_style=False)
+            with open(self.filename, 'w') as outfile:
+                dump_yaml(out, outfile, default_flow_style=False)
         except Exception as e:
-            raise PluginError(f'Error writhing data to `{self._yaml_file}`: {e}')
-
-    def get(self, item) -> None:
-        title = item.get('title', None)
-        if not title:
-            logger.error('Can\'t get entry, no `title` field')
-            return
+            raise PluginError(f'Error writhing data to `{self.filename}`: {e}')
 
-        item = next((data for key, data in self._yaml_items.items() if key == title), None)
-        if not item:
-            return None
+    def get(self, item) -> Optional[Entry]:
+        for entry in self.entries:
+            if self.matches(item, entry):
+                return entry
+        return None
 
-        item = self._output_item(title, item)
-        return item
-
-    def add(self, item) -> None:
-        title = item.get('title', None)
-        if not title:
-            logger.error('Can\'t add entry, no `title` field')
+    def add(self, entry: Entry) -> None:
+        exists = self.get(entry)
+        if exists:
             return
-
-        if isinstance(item, Entry):
-            new_item_dict = {}
-            for key, data in item.items():
-                new_item_dict[key] = data
-            item = new_item_dict
-
-        item_dict = self._get_item_fields(item)
-        item_dict = {title: jsonify(item_dict)}
-        item_dict.pop('title', None)
-        self._yaml_items.update(item_dict)
+        self.entries.append(entry)
         self.save_yaml()
 
     def discard(self, item) -> None:
@@ -211,12 +110,18 @@ def discard(self, item) -> None:
             logger.error('Can\'t add entry, no `title` field')
             return
 
-        self._yaml_items.pop(title, None)
+        for i, entry in enumerate(self.entries):
+            if self.matches(item, entry):
+                self.entries.pop(i)
+                break
+        else:
+            return
+
         self.save_yaml()
 
     @property
     def online(self):
-        return True
+        return False
 
     @property
     def immutable(self):
@@ -240,41 +145,10 @@ class YamlList:
     }
 
     def process_config(self, config: dict) -> dict:
-        """Process config
-
-        Args:
-            config (dict): Config to process
-
-        Returns:
-            dict: Processed config
-        """
-
-        new_config = {}
         if isinstance(config, str):
-            new_config['path'] = config
-        else:
-            new_config = config
-
-        if 'fields' not in new_config:
-            new_config['fields'] = []
-
-        return new_config
-
-    def item_to_entry(self, item: dict) -> Entry:
-        """Get's entry from item
-
-        Args:
-            item (dict): Item to return as entry
-
-        Returns:
-            Entry: Entry
-        """
-
-        if not item:
-            return None
-        new_entry = Entry()
-        new_entry.update(item)
-        return new_entry
+            config = {'path': config}
+        config.setdefault('fields', [])
+        return config
 
     def get_list(self, config):
         config = self.process_config(config)
@@ -285,7 +159,7 @@ def on_task_input(self, task, config):
         config = self.process_config(config)
         yaml_list = YamlManagedList(**config)
         for item in yaml_list:
-            yield self.item_to_entry(item)
+            yield item
 
 
 @event('plugin.register')