112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
import json
|
|
|
|
from modules.civit.fetch import Fetch
|
|
from shell.Handlers.ABS import Handler
|
|
from modules.civit.client import Client
|
|
|
|
class CivitHandler(Handler):
|
|
|
|
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.forwarding_table: dict[str, Handler] = {
|
|
'fetch': FetchHandler(self)
|
|
}
|
|
self.handle_table: dict = {
|
|
'init': self._init,
|
|
}
|
|
self.client: Client | None = None
|
|
|
|
|
|
def _init(self, command: list[str], pos=0):
|
|
keys, args = self.parse_arguments(command[pos:], ['path', 'key'])
|
|
self._check_arg(keys, 'path')
|
|
self.client = Client(keys['path'], keys['key'])
|
|
self.succeed = True
|
|
|
|
|
|
class FetchHandler(Handler):
|
|
|
|
|
|
|
|
def __init__(self, parent):
|
|
super().__init__()
|
|
self.parent: CivitHandler = parent
|
|
self.forwarding_table: dict[str, Handler] = {
|
|
}
|
|
self.handle_table: dict = {
|
|
'creators_raw': self._creators_raw,
|
|
'creators': self._creators,
|
|
'tags_raw': self._tags_raw,
|
|
'tags': self._tags,
|
|
'models': self._models,
|
|
'images': self._images,
|
|
'datamodel': self._datamodel,
|
|
'load': self._load,
|
|
|
|
}
|
|
|
|
def _load(self, command: list[str], pos=0):
|
|
keys, args = self.parse_arguments(command[pos:], ['entity'])
|
|
self._check_arg(keys, 'entity')
|
|
|
|
res = Fetch.load(self.parent.client, keys['entity'])
|
|
for r in res: print(r)
|
|
self.succeed = True
|
|
|
|
|
|
|
|
def _creators_raw(self, command: list[str], pos=0):
|
|
keys, args = self.parse_arguments(command[pos:], ['page', 'limit', 'query'])
|
|
res = self.parent.client.get_creators_raw(page=keys['page'], limit=keys['limit'], query=keys['query'])
|
|
print(res)
|
|
self.succeed = True
|
|
|
|
def _creators(self, command: list[str], pos=0):
|
|
res = Fetch.creators(self.parent.client)
|
|
for r in res: print(r)
|
|
self.succeed = True
|
|
|
|
|
|
def _tags_raw(self, command: list[str], pos=0):
|
|
keys, args = self.parse_arguments(command[pos:], ['page', 'limit', 'query'])
|
|
res = self.parent.client.get_tags_raw(page=keys['page'], limit=keys['limit'], query=keys['query'])
|
|
print(res)
|
|
self.succeed = True
|
|
|
|
def _tags(self, command: list[str], pos=0):
|
|
res = Fetch.tags(self.parent.client)
|
|
for r in res: print(r)
|
|
self.succeed = True
|
|
|
|
def _models(self, command: list[str], pos=0):
|
|
res = Fetch.models(self.parent.client)
|
|
for r in res: print(r)
|
|
self.succeed = True
|
|
|
|
def _images(self, command: list[str], pos=0):
|
|
keys, args = self.parse_arguments(command[pos:], ['start'])
|
|
res = Fetch.images(self.parent.client, start_with=keys['start'])
|
|
for r in res: print(r)
|
|
self.succeed = True
|
|
|
|
def _datamodel(self, command: list[str], pos=0):
|
|
entities = {
|
|
'creator': 'fetch_creators',
|
|
'tag': 'fetch_tags',
|
|
'model': 'fetch_models',
|
|
'image': 'fetch_images',
|
|
}
|
|
keys, args = self.parse_arguments(command[pos:], ['entity', 'top', 'dump'])
|
|
self._check_arg(keys, 'entity')
|
|
if keys['entity'] in entities:
|
|
res = Fetch.datamodel(self.parent.client, entities[keys['entity']], keys['top'] or 10)
|
|
print(json.dumps(res, indent=2, ensure_ascii=False))
|
|
if keys['dump']:
|
|
with open(keys['dump'], 'w') as f:
|
|
json.dump(res, f, indent=2, ensure_ascii=False)
|
|
|
|
|
|
self.succeed = True
|