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
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
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):
|
||||
@ -57,7 +201,7 @@ def main():
|
||||
# Flag if the bot is running in debug mode
|
||||
debug_mode = arguments.dry_run or arguments.debug
|
||||
|
||||
yandere = yandere_bot.YandereBot(
|
||||
yandere = YandereBot(
|
||||
yandere_config,
|
||||
arguments.keyfile,
|
||||
debug_mode
|
||||
@ -84,9 +228,9 @@ if __name__ == "__main__":
|
||||
except FailedToLoadCfg:
|
||||
sys.exit(10)
|
||||
# Exceptions raised from the bot
|
||||
except yandere_bot.Debug:
|
||||
except FediBot.Debug:
|
||||
sys.exit(6)
|
||||
except yandere_bot.BadCfgFile:
|
||||
except FediBot.BadCfgFile:
|
||||
sys.exit(4)
|
||||
except yandere_bot.FailedLogin:
|
||||
except FediBot.FailedLogin:
|
||||
sys.exit(2)
|
||||
|
Reference in New Issue
Block a user