Updated to make project modular
This commit is contained in:
parent
98cc739c67
commit
fe00f1d635
154
src/main.py
154
src/main.py
@ -16,10 +16,154 @@
|
|||||||
# You should have received a copy of the GNU General Public License
|
# You should have received a copy of the GNU General Public License
|
||||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||||
|
|
||||||
|
import os
|
||||||
import sys
|
import sys
|
||||||
import argparse
|
import argparse
|
||||||
import signal
|
import signal
|
||||||
import yandere_bot
|
import random
|
||||||
|
import datetime
|
||||||
|
import subprocess
|
||||||
|
import FediBot
|
||||||
|
from mastodon import MastodonAPIError
|
||||||
|
|
||||||
|
# A class that contains all of the rendering information for a single post
|
||||||
|
class PickAndRenderFrame:
|
||||||
|
# From Config File
|
||||||
|
profile_name = ""
|
||||||
|
path = ""
|
||||||
|
frames = 0
|
||||||
|
skip = tuple()
|
||||||
|
nsfw = tuple()
|
||||||
|
output_name = ""
|
||||||
|
message = ""
|
||||||
|
message_nsfw = ""
|
||||||
|
render_script = ""
|
||||||
|
|
||||||
|
# After pick
|
||||||
|
picked_frame = None
|
||||||
|
output_name_tr = ""
|
||||||
|
|
||||||
|
def __init__(self, picked, datetime_format):
|
||||||
|
dt_now = datetime.datetime.now()
|
||||||
|
dt_now_str = datetime.datetime.strftime(dt_now, datetime_format)
|
||||||
|
self.profile_name = picked["profile_name"]
|
||||||
|
self.path = picked["path"]
|
||||||
|
self.frames = picked["frames"]
|
||||||
|
self.skip = picked["skip"]
|
||||||
|
self.nsfw = picked["nsfw"]
|
||||||
|
self.output_name = picked["output_name"]
|
||||||
|
self.output_name_tr = self.output_name
|
||||||
|
self.message = picked["message"]
|
||||||
|
self.message_nsfw = picked["message_nsfw"]
|
||||||
|
self.render_script = picked["render_script"]
|
||||||
|
|
||||||
|
self.picked_frame = self._pick_frame()
|
||||||
|
|
||||||
|
# Shell-like substitutions
|
||||||
|
translations = {
|
||||||
|
"profile_name": self.profile_name,
|
||||||
|
"frame": str(self.picked_frame),
|
||||||
|
"datetime": dt_now_str
|
||||||
|
}
|
||||||
|
|
||||||
|
self.output_name_tr = self._translate_basename(translations)
|
||||||
|
self._render_frame()
|
||||||
|
|
||||||
|
def __del__(self):
|
||||||
|
render_file = self.output_name_tr
|
||||||
|
if render_file and os.path.isfile(render_file):
|
||||||
|
os.remove(render_file)
|
||||||
|
|
||||||
|
def _pick_frame(self):
|
||||||
|
picked_frame = None
|
||||||
|
while picked_frame is None:
|
||||||
|
picked_frame = random.random() * self.frames
|
||||||
|
for skip in self.skip:
|
||||||
|
begin, end = skip
|
||||||
|
if begin <= picked_frame <= end:
|
||||||
|
picked_frame = None
|
||||||
|
break
|
||||||
|
return picked_frame
|
||||||
|
|
||||||
|
def is_nsfw(self):
|
||||||
|
for nsfw in self.nsfw:
|
||||||
|
begin, end = nsfw
|
||||||
|
if begin <= self.picked_frame <= end:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _translate_basename(self, translations):
|
||||||
|
output_name_tr = self.output_name
|
||||||
|
for k,v in translations.items():
|
||||||
|
replace_token = "${" + k + "}"
|
||||||
|
output_name_tr = output_name_tr.replace(replace_token, v)
|
||||||
|
return output_name_tr
|
||||||
|
|
||||||
|
# TODO: Add translation keywords to the message
|
||||||
|
def get_message(self):
|
||||||
|
return self.message_nsfw if self.is_nsfw() else self.message
|
||||||
|
|
||||||
|
def _render_frame(self):
|
||||||
|
args = [
|
||||||
|
self.render_script,
|
||||||
|
self.path,
|
||||||
|
self.output_name_tr,
|
||||||
|
str(self.picked_frame)
|
||||||
|
]
|
||||||
|
subprocess.run(args)
|
||||||
|
|
||||||
|
|
||||||
|
class YandereBot(FediBot.YandereBot):
|
||||||
|
def __init__(self, cfg, keyfile=None, debug_mode=False):
|
||||||
|
settings = {
|
||||||
|
"settings_time": {},
|
||||||
|
"settings_post": {},
|
||||||
|
}
|
||||||
|
self.settings.update(settings)
|
||||||
|
super(YandereBot, self).__init__(cfg, keyfile, debug_mode)
|
||||||
|
|
||||||
|
# Maybe I should remove this from the backend?
|
||||||
|
def print_header_stats(self, picked):
|
||||||
|
profile, frame, nsfw, path = None, None, None, None
|
||||||
|
if picked:
|
||||||
|
profile = picked.profile_name
|
||||||
|
frame = picked.picked_frame
|
||||||
|
nsfw = picked.is_nsfw()
|
||||||
|
path = picked.output_name_tr
|
||||||
|
print("Profile: {} | Frame: {} | NSFW: {} | Path: {}".format(
|
||||||
|
profile, frame, nsfw, path
|
||||||
|
))
|
||||||
|
|
||||||
|
|
||||||
|
def pick(self):
|
||||||
|
dt_picked = self.settings["settings_time"]["datetime"]
|
||||||
|
picked_profile = random.choice(self.settings["settings_post"])
|
||||||
|
|
||||||
|
picked = PickAndRenderFrame(picked_profile, dt_picked)
|
||||||
|
media_list = [picked.output_name_tr]
|
||||||
|
spoiler = picked.is_nsfw()
|
||||||
|
message = picked.get_message()
|
||||||
|
return {
|
||||||
|
"picked": picked,
|
||||||
|
"media_list": media_list,
|
||||||
|
"spoiler": spoiler,
|
||||||
|
"message": message
|
||||||
|
}
|
||||||
|
|
||||||
|
def after_pick(self, picked):
|
||||||
|
self.print_header_stats(picked["picked"])
|
||||||
|
|
||||||
|
def post(self, picked):
|
||||||
|
try:
|
||||||
|
super(YandereBot, self).post(picked)
|
||||||
|
except (FileNotFoundError, MastodonAPIError, FediBot.InvalidPost, Exception) as e:
|
||||||
|
print("Exception:", e)
|
||||||
|
|
||||||
|
self.handle_post_exception()
|
||||||
|
|
||||||
|
# The post failed
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
class FailedToLoadCfg(Exception):
|
class FailedToLoadCfg(Exception):
|
||||||
@ -57,7 +201,7 @@ def main():
|
|||||||
# Flag if the bot is running in debug mode
|
# Flag if the bot is running in debug mode
|
||||||
debug_mode = arguments.dry_run or arguments.debug
|
debug_mode = arguments.dry_run or arguments.debug
|
||||||
|
|
||||||
yandere = yandere_bot.YandereBot(
|
yandere = YandereBot(
|
||||||
yandere_config,
|
yandere_config,
|
||||||
arguments.keyfile,
|
arguments.keyfile,
|
||||||
debug_mode
|
debug_mode
|
||||||
@ -84,9 +228,9 @@ if __name__ == "__main__":
|
|||||||
except FailedToLoadCfg:
|
except FailedToLoadCfg:
|
||||||
sys.exit(10)
|
sys.exit(10)
|
||||||
# Exceptions raised from the bot
|
# Exceptions raised from the bot
|
||||||
except yandere_bot.Debug:
|
except FediBot.Debug:
|
||||||
sys.exit(6)
|
sys.exit(6)
|
||||||
except yandere_bot.BadCfgFile:
|
except FediBot.BadCfgFile:
|
||||||
sys.exit(4)
|
sys.exit(4)
|
||||||
except yandere_bot.FailedLogin:
|
except FediBot.FailedLogin:
|
||||||
sys.exit(2)
|
sys.exit(2)
|
||||||
|
Reference in New Issue
Block a user