Add /code feature.

2025-08-21 15:27:14 +00:00 · 2021-01-08 17:32:39 +03:00
parent 8f3cccf01d
commit 4b142ae8ef
8 changed files with 916 additions and 771 deletions
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "desubot"
-version = "0.1.0"
+version = "0.5.0"
 authors = ["AB <ab@hexor.ru>"]
 edition = "2018"

@@ -33,3 +33,8 @@ rand = "0.7.3"
 mystem = "0.2.1"
 async-trait = "0.1.42"
 sqlparser = "0.7.0"
+
+[dependencies.syntect]
+version = "4.4"
+default-features = false
+features = ["parsing", "dump-load", "regex-onig"]
--- a/2
+++ b/2
@@ -9,7 +9,7 @@ Telegram bot with light group statistic and heavy spy features.
 * Generate sentences using Markov Chains trained on history with /markov_all.

 == TODO ==
-* Syntax highlighting for code exported to image.
+* Syntax highlighting for CODE exported to image.

 == Important ==
 * Desubot uses MyStem by Yandex for word stemming and assume that mystem binary is available in PATH.
--- a/assets/help_text.rs
+++ b/assets/help_text.rs
@@ -0,0 +1,36 @@
+static CODE: &str = "<b>Code highlighter</b>
+
+<i>Usage</i>
+<pre>/CODE
+#&lt;theme - Dracula by default&gt;
+&lt;CODE&gt;
+#&lt;lang - JS by default&gt;</pre>
+
+Language may be defined by both name and extension - Rust, rs...
+Max lines - 80
+
+List of themes:
+1337
+DarkNeon
+Dracula
+GitHub
+Monokai_Extended
+Monokai_Extended_Bright
+Monokai_Extended_Light
+Monokai_Extended_Origin
+Nord
+OneHalfDark
+OneHalfLight
+Solarized_(dark)
+Solarized_(light)
+Sublime_Snazzy
+TwoDark
+ansi-dark
+ansi-light
+base16
+base16-256
+gruvbox
+gruvbox-light
+gruvbox-white
+zenburn
+";
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -1,7 +1,7 @@
 #![allow(unused_variables)]
 use crate::db;
 use crate::errors::Error;
-use crate::errors::Error::{SQLITE3Error, SQLInvalidCommand};
+use crate::errors::Error::{CodeHighlightningError, SQLITE3Error, SQLInvalidCommand};
 use async_trait::async_trait;
 use html_escape::encode_text;
 use markov::Chain;
@@ -13,12 +13,15 @@ use mystem::Tense::{Inpresent, Past};
 use rand::seq::SliceRandom;
 use rand::Rng;
 use regex::Regex;
-
 use sqlparser::ast::Statement;
 use sqlparser::dialect::GenericDialect;
 use sqlparser::parser::Parser;
+use syntect::easy::HighlightLines;
+use syntect::parsing::SyntaxReference;
+use syntect::util::LinesWithEndings;
 use telegram_bot::prelude::*;
 use telegram_bot::{Api, Message, ParseMode};
+use syntect::highlighting::Theme;

 pub struct Here {
    pub data: String,
@@ -38,7 +41,6 @@ pub struct Omedeto {
 pub struct Sql {
    pub data: String,
 }
-
 pub struct Code {
    pub data: String,
 }
@@ -396,6 +398,7 @@ impl Execute for Omedeto {
            .map(|m| m.split(' ').map(|s| s.to_string()).collect::<Vec<String>>()[1].clone())
            .filter(|m| {
                let stem = mystem.stemming(m.clone()).unwrap_or_default();
+                #[allow(clippy::if_same_then_else)]
                if stem.is_empty() {
                    false
                } else if stem[0].lex.is_empty() {
@@ -423,6 +426,7 @@ impl Execute for Omedeto {
            .map(|m| m.split(' ').map(|s| s.to_string()).collect::<Vec<String>>()[1].clone())
            .filter(|m| {
                let stem = mystem.stemming(m.clone()).unwrap_or_default();
+                #[allow(clippy::if_same_then_else)]
                if stem.is_empty() {
                    false
                } else if stem[0].lex.is_empty() {
@@ -478,13 +482,11 @@ impl Execute for Omedeto {
                .map(|m| m.split(' ').map(|s| s.to_string()).collect::<Vec<String>>()[1].clone())
                .map(|m| {
                    let stem = mystem.stemming(m).unwrap_or_default();
+                    #[allow(clippy::if_same_then_else)]
                    if stem.is_empty() {
-
                    } else if stem[0].lex.is_empty() {
-
                    } else {
-                        match stem[0].lex[0].grammem.part_of_speech {
-                            mystem::PartOfSpeech::Verb => {
+                        if let mystem::PartOfSpeech::Verb = stem[0].lex[0].grammem.part_of_speech {
                            match stem[0].lex[0]
                                .grammem
                                .facts
@@ -504,8 +506,6 @@ impl Execute for Omedeto {
                                false => (),
                            }
                        }
-                            _ => (),
-                        }
                    }
                })
                .for_each(drop);
@@ -518,17 +518,17 @@ impl Execute for Omedeto {
            start.choose(&mut rand::thread_rng()).unwrap(),
            message.from.first_name.to_string(),
            { if fem { "ая" } else { "ый" } },
-            nouns.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
-            nouns.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
-            nouns.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
+            nouns.pop().unwrap_or_else(|| placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
+            nouns.pop().unwrap_or_else(|| placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
+            nouns.pop().unwrap_or_else(|| placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
            { if fem { "а" } else { "" } },
-            verbs_p.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
-            verbs_p.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
-            verbs_p.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
+            verbs_p.pop().unwrap_or_else(|| placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
+            verbs_p.pop().unwrap_or_else(|| placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
+            verbs_p.pop().unwrap_or_else(|| placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
            { if fem { "а" } else { "" } },
-            verbs_i.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
-            verbs_i.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
-            verbs_i.pop().unwrap_or(placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
+            verbs_i.pop().unwrap_or_else(|| placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
+            verbs_i.pop().unwrap_or_else(|| placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
+            verbs_i.pop().unwrap_or_else(|| placeholders.choose(&mut rand::thread_rng()).unwrap().to_string()),
        );
        match api
            .send(
@@ -548,21 +548,81 @@ impl Execute for Omedeto {
 #[async_trait]
 impl Execute for Code {
    async fn exec(&self, api: &Api, message: &Message) -> Result<(), Error> {
-        let lang = "Rust";
-        let theme = "Dracula";
-        let code = &self.data;
-        let (ps, ts) = silicon::utils::init_syntect();
-        let syntax = ps
-            .find_syntax_by_token(lang)
-            .ok_or_else(|| ps.find_syntax_by_token("js"));
-        let theme = ts.themes.get(theme).ok_or_else(|| ts.themes.get("1337"));
-        debug!("{:?}", code);
-        Ok(())
-    }
-
-    async fn exec_with_result(&self, api: &Api, message: &Message) -> Result<String, Error> {
        unimplemented!()
    }
+    async fn exec_with_result(&self, api: &Api, message: &Message) -> Result<String, Error> {
+        let mut lines: Vec<String> = self.data.trim().split("\n").map(|s| s.to_string()).collect();
+        if lines.len() >= 81 {
+            return Err(CodeHighlightningError);
+        }
+        let last_line = &lines[lines.len()-1];
+
+        let tags = last_line
+            .trim()
+            .split(|s| s == ' ' || s == '\n')
+            .filter(|s| s.starts_with("#"))
+            .map(|s| s.to_string().replace("#", ""))
+            .map(|s| s.to_string().replace("_", " "))
+            .collect::<Vec<_>>();
+
+        let code = if tags.is_empty() {
+            self.data.trim().to_string()
+        } else {
+            let _ = lines.pop();
+            lines.join("\n")
+        };
+        if code.is_empty() {
+            return Err(CodeHighlightningError);
+        }
+        let (ps, ts) = silicon::utils::init_syntect();
+        let syntax: Vec<&SyntaxReference> = tags
+            .iter()
+            .map(|s| ps.find_syntax_by_token(s))
+            .filter(|s| s.is_some())
+            .map(|s| s.unwrap())
+            .collect();
+        let syntax = if syntax.len() != 1 {
+            ps.find_syntax_by_token("js").unwrap()
+        } else {
+            syntax[0]
+        };
+        let theme: Vec<&Theme> = tags
+            .iter()
+            .map(|s| ts.themes.get(s))
+            .filter(|s| s.is_some())
+            .map(|s| s.unwrap())
+            .collect();
+
+
+        let theme = if theme.len() != 1 {
+            ts.themes.get("Dracula").unwrap()
+        } else {
+            theme[0]
+        };
+
+        let mut h = HighlightLines::new(syntax, theme);
+        let highlight = LinesWithEndings::from(&code)
+            .map(|line| h.highlight(line, &ps))
+            .collect::<Vec<_>>();
+        let formatter = silicon::formatter::ImageFormatterBuilder::<String>::new()
+            .window_controls(false)
+            .round_corner(false);
+        let mut formatter = formatter.build().unwrap();
+        let image = formatter.format(&highlight, &theme);
+        let path = "code.png";
+        image
+            .save(&path)
+            .map_err(|e| error!("Failed to save image to {}: {}", path, e))
+            .unwrap();
+
+        // let file = InputFileUpload::with_path("CODE.png");
+        // api.send(message.chat.document(&file)).await?;
+        //
+        // // Send an image from disk
+        // api.send( message.chat.photo(&file)).await?;
+        //debug!("{:#?}", formatter);
+        Ok(path.into())
+    }

    async fn exec_mystem(
        &self,
--- a/src/db.rs
+++ b/src/db.rs
@@ -348,10 +348,9 @@ pub(crate) async fn get_file(file_id: String) -> Result<i64, errors::Error> {
    file_rowid
 }

-async fn add_word(word: &String) -> Result<i64, errors::Error> {
-    match get_stop_word(&word).await {
-        Err(_) => return Err(errors::Error::WordInStopList),
-        _ => {}
+async fn add_word(word: &str) -> Result<i64, errors::Error> {
+    if get_stop_word(&word).await.is_err() {
+        return Err(errors::Error::WordInStopList);
    }
    let conn = open()?;
    let word_rowid =
@@ -365,7 +364,7 @@ async fn add_word(word: &String) -> Result<i64, errors::Error> {
    Ok(word_rowid)
 }

-async fn get_stop_word(stop_word: &String) -> Result<(), errors::Error> {
+async fn get_stop_word(stop_word: &str) -> Result<(), errors::Error> {
    let conn = open()?;
    match conn.execute_named(
        "SELECT rowid FROM stop_words WHERE word = (:stop_word)",
--- a/src/errors.rs
+++ b/src/errors.rs
@@ -23,11 +23,39 @@ pub enum Error {
    SQLBannedCommand(String),
    SQLInvalidCommand,
    SQLResultTooLong(String),
+    CodeHighlightningError,
 }

 impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "An error occurred.")
+        //         match self {
+        //             _ => write!(f, "An error occurred."),
+        //             // Error::UserNotFound => {}
+        //             // Error::SQLITE3Error(_) => {}
+        //             // Error::TelegramError(_) => {}
+        //             // Error::ReqwestError(_) => {}
+        //             // Error::ConfNotFound => {}
+        //             // Error::WordNotFound => {}
+        //             // Error::WordInStopList => {}
+        //             // Error::IOError(_) => {}
+        //             // Error::FileNotFound => {}
+        //             // Error::JsonParseError(_) => {}
+        //             // Error::PopenError(_) => {}
+        //             // Error::MystemError(_) => {}
+        //             // Error::SQLBannedCommand(_) => {}
+        //             // Error::SQLInvalidCommand => {}
+        //             // Error::SQLResultTooLong(_) => {}
+        // //             Error::CodeHighlightningError(Help) => write!(f, "Code highlighter.\
+        // // <b>Usage</b><pre>/CODE\
+        // // #&lt;theme&gt;\
+        // // &lt;CODE&gt;\
+        // // #&lt;lang&gt;</pre>\
+        // // \
+        // // List of themes:\
+        // // .")
+        //             Error::CodeHighlightningError(help) => write!(f, "{}", help.description)
+        //         }
    }
 }

--- a/src/handlers.rs
+++ b/src/handlers.rs
@@ -6,6 +6,9 @@ use crate::utils;
 use mystem::MyStem;
 use telegram_bot::*;

+
+include!("../assets/help_text.rs");
+
 pub async fn handler(
    api: Api,
    message: Message,
@@ -22,12 +25,38 @@ pub async fn handler(
                title,
                &message.from.id,
                &message.from.first_name,
-                data
+                {if data.len() <= 200 {data} else {&data[..200]}}.replace("\n", " ")
            );
-            db::add_sentence(&message, mystem).await?;
+
            let cleaned_message = data.replace(&format!("@{}", me.clone().username.unwrap()), "");
            match cleaned_message.as_str() {
+                s if s.to_string().starts_with("/code") => {
+
+                    match {
+                        Code {
+                            data: s.replace("/code", ""),
+                        }
+                            .exec_with_result(&api, &message)
+                            .await
+                    } {
+                        Ok(path) => {
+                            let file = InputFileUpload::with_path(path.clone());
+                            // api.send(message.chat.document(&file)).await?;
+                            //
+                            // // Send an image from disk
+                            api.send(message.chat.photo(&file)).await?;
+                            //debug!("{:#?}", formatter);
+                            let _ = std::fs::remove_file(&path);
+                        }
+                        Err(_) => {
+                            let _ = api
+                                .send(message.text_reply(CODE).parse_mode(ParseMode::Html))
+                                .await?;
+                        }
+                    }
+                }
                s if s.contains("/here") => {
+                    db::add_sentence(&message, mystem).await?;
                    Here {
                        data: "".to_string(),
                    }
@@ -43,13 +72,9 @@ pub async fn handler(
                } {
                    Ok(msg) => {
                        let _ = api
-                            .send(
-                                message
-                                    .text_reply(msg)
-                                    .parse_mode(ParseMode::Html),
-                            )
+                            .send(message.text_reply(msg).parse_mode(ParseMode::Html))
                            .await?;
-                    },
+                    }
                    Err(e) => {
                        let _ = api
                            .send(
@@ -60,13 +85,6 @@ pub async fn handler(
                            .await?;
                    }
                },
-                s if s.to_string().starts_with("/code") => {
-                    Code {
-                        data: s.to_string(),
-                    }
-                        .exec(&api, &message)
-                        .await?
-                }
                "/top" => {
                    Top {
                        data: "".to_string(),
@@ -102,7 +120,7 @@ pub async fn handler(
                    .exec_mystem(&api, &message, mystem)
                    .await?
                }
-                _ => (),
+                _ => db::add_sentence(&message, mystem).await?,
            }
        }
        MessageKind::Photo { ref caption, .. } => {
@@ -113,7 +131,7 @@ pub async fn handler(
                title,
                &message.from.id,
                &message.from.first_name,
-                caption.clone().unwrap_or("NO_TITLE".to_string())
+                caption.clone().unwrap_or_else(|| "NO_TITLE".to_string())
            );
            utils::get_files(api, message, token).await?;
        }
@@ -126,7 +144,7 @@ pub async fn handler(
                title,
                &message.from.id,
                &message.from.first_name,
-                caption.clone().unwrap_or("NO_TITLE".to_string())
+                caption.clone().unwrap_or_else(|| "NO_TITLE".to_string())
            );
            utils::get_files(api, message, token).await?;
        }
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -18,10 +18,9 @@ pub(crate) fn get_title(message: &Message) -> String {
    }
 }

-pub(crate) async fn create_dir(dir: &String)  {
-    match fs_create_dir(dir) {
-        Ok(_) => info!("Dir {} created.", dir),
-        Err(_) => (),
+pub(crate) async fn create_dir(dir: &str) {
+    if fs_create_dir(dir).is_ok() {
+        info!("Dir {} created.", dir)
    }
 }