Krita/krita/pykrita/sketchbookmode/sketchbookmode.py

import random
import time
import os
from pathlib import Path
from datetime import datetime, timedelta
from human_id import generate_id
from krita import Krita, Selection, Extension, ManagedColor
from PyQt5.QtGui import QColor, QIcon, QPixmap
from PyQt5.QtCore import QByteArray

dt_fmt = "%Y-%m-%d-%H-%M-%S"
ccolors = ['#ffb482', '#a1c9f4', '#8de5a1', '#ff9f9b', '#d0bbff', '#debb9b', '#fab0e4', '#fffea3', '#b9f2f0']
history = {}
FOLDERS = "grey studies animations sfw pieces sketches lineart doodles ideas comics lo".split()

def fill_color(app, doc, view, x, y, w, h, color):
  newSelection = Selection()
  opacity = 255
  newSelection.select(x, y, w, h, opacity)
  doc.setSelection(newSelection)
  mColor = ManagedColor.fromQColor(color)
  view.setForeGroundColor(mColor)
  doc.waitForDone()
  app.action('fill_selection_foreground_color').trigger()
  doc.waitForDone()
  mColor = ManagedColor.fromQColor(QColor(0, 0, 0))
  view.setForeGroundColor(mColor)

#TODO: write date on page using text tool
def write_text(app, doc, view, x, y, w, h, text):
  newSelection = Selection()
  opacity = 255
  newSelection.select(x, y, w, h, opacity)
  doc.setSelection(newSelection)
  app.action('SvgTextTool').trigger()
  doc.waitForDone()
  mColor = ManagedColor.fromQColor(QColor(0, 0, 0))
  view.setForeGroundColor(mColor)


def create_new_page2(path: Path):
  print(path)
  app = Krita.instance()
  win = app.activeWindow()
  view = win.activeView()
  base = path.joinpath("base.kra.tmpl").as_posix()
  now = datetime.now()
  w, h = 5000, 5000
  new = Path(path).joinpath(f"{now.strftime(dt_fmt)}.kra").as_posix()
  os.system(f"""cp '{base}' '{new}'""")
  doc = app.openDocument(new)
  app.activeWindow().addView(doc)

  random.seed(time.time())
  square_size = 200
  randcolor2 = QColor(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
  black = QColor(35, 7, 33)
  border = 15
  fill_color(app, doc, view, w-square_size-100-border, h-square_size-100-border, square_size+(border*2), square_size+(border*2), black)
  fill_color(app, doc, view, w-square_size-100, h-square_size-100, square_size, square_size, randcolor2)

  app.action('merge_layer').trigger()
  doc.waitForDone()
  app.action('deselect').trigger()
  doc.waitForDone()
  newLayer2 = doc.createNode("PaintLayer", "paintlayer")
  doc.rootNode().addChildNode(newLayer2, None)
  doc.save()
  app.action('deselect').trigger()
  doc.waitForDone()


# def create_new_page(path: Path):
#   raise Exception("This function should never be called.")
#   app = Krita.instance()
#   win = app.activeWindow()
#   view = win.activeView()
#   w = 3840
#   h = 2160
#   square_size = 200
#   now = datetime.now()
#   now_str = now.strftime(dt_fmt)
#   colorspace = "GRAYA" if path.name in ["grey", "lineart"] else "RGBA"
#   doc = app.createDocument(w, h, now_str, colorspace, "U8", "", 100.0)
#   v = win.addView(doc)
#   newLayer = doc.createNode("PaintLayer", "paintlayer")
#   doc.rootNode().addChildNode(newLayer, None)
#   doc.setAnnotation(f'sketchbook_page', f"Sketchbook Page Number", QByteArray(now_str.encode()))
#   doc.waitForDone()
#   doc.setFileName(path.joinpath(f"{now_str}.kra").as_posix())
#   doc.setName(generate_id(seed=now_str, separator='-', word_count=3))
#   doc.waitForDone()
#
#   random.seed(time.time())
#   if path.name == "sketches":
#     t = int(w/3)
#     fill_color(app, doc, view, 0  , 0,   t, h, QColor(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
#     fill_color(app, doc, view, t  , 0, 2*t, h, QColor(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
#     fill_color(app, doc, view, 2*t, 0, 3*t, h, QColor(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
#   elif path.name == "lineart":
#     fill_color(app, doc, view, 0, 0, w, h, QColor(100, 100, 100))
#   elif path.name == "doodles":
#     fill_color(app, doc, view, 0, 0, w, h, QColor(130, 130, 130))
#   else:
#     fill_color(app, doc, view, 0, 0, w, h, QColor(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
#
#   random.seed(path.name)
#   randcolor2 = QColor(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
#   black = QColor(35, 7, 33)
#   border = 15
#   fill_color(app, doc, view, w-square_size-100-border, h-square_size-100-border, square_size+(border*2), square_size+(border*2), black)
#   fill_color(app, doc, view, w-square_size-100, h-square_size-100, square_size, square_size, randcolor2)
#
#   app.action('merge_layer').trigger()
#   doc.waitForDone()
#   app.action('deselect').trigger()
#   doc.waitForDone()
#   newLayer2 = doc.createNode("PaintLayer", "paintlayer")
#   doc.rootNode().addChildNode(newLayer2, None)
#   doc.save()
#   win.activate()
#   win.showView(v)
#   v.setVisible()
#   currentPreset = v.currentBrushPreset()
#   v.showFloatingMessage(
#     generate_id(seed=now_str, separator='-', word_count=3),
#     QIcon(QPixmap.fromImage(currentPreset.image())),
#     8000,
#     99
#   )
#
#
# def get_last_page(path: Path) -> None:
#   raise Exception("This function should never be called.")
#   app = Krita.instance()
#   win = app.activeWindow()
#   view = win.activeView()
#   paths = path.glob("*.kra")
#   next_page_candidate = None
#   i = 0
#   l = sorted(paths, key=lambda w: w.as_posix())
#   for p in l:
#     i += 1
#     try:
#       dt = datetime.strptime(Path(p).name, f"{dt_fmt}.kra")
#       if next_page_candidate is None or dt > next_page_candidate:
#         next_page_candidate = dt
#     except Exception as e:
#       print(e)
#   if next_page_candidate is not None:
#     #Check if next page already open in krita
#     for v in app.activeWindow().views():
#       if v.document().fileName() == next_page_candidate.strftime(dt_fmt) + ".kra":
#         print(v.document().annotation("sketchbook_page"))
#     b = Path(path).joinpath(next_page_candidate.strftime(dt_fmt) + ".kra").as_posix()
#     doc2 = app.openDocument(b)
#     app.activeWindow().addView(doc2)
#     currentPreset = v.currentBrushPreset()
#     o = f"{Path(path).stem.upper()} "
#     o += generate_id(seed=next_page_candidate.strftime(dt_fmt), separator='-', word_count=3)
#     o += f" {i}/{len(l)}"
#     v.showFloatingMessage(
#       o,
#       QIcon(QPixmap.fromImage(currentPreset.image())),
#       5000,
#       99
#     )
#   else:
#     print("Failure")
#     # create_new_page(path)
#   # TODO: automatically close next and previous pages +3 to save memory

def animation_tools_toggle(sketchbookname):
  app = Krita.instance()
  win = app.activeWindow()
  doc =  app.activeDocument()

  for w in win.dockers():
    # print(w.objectName())
    # if w.objectName() in ["TimelineDocker", "OnionSkinsDocker", "AnimationCurvesDocker"]:
    if w.objectName() in ["TimelineDocker"]:
      print('>a', sketchbookname in ['animations'])
      print('>b', bool(doc.animationLength() > 1))
      print('>c', doc.animationLength())
      w.setVisible(sketchbookname in ['animations'] or doc.animationLength() > 1)
      # print(w.isEnabled())

def get_next_prev_page(path: Path, current_page_dt: datetime, prev=False, exact=False) -> None:
  app = Krita.instance()
  doc = app.activeDocument()
  win = app.activeWindow()
  view = win.activeView()
  paths = path.glob("*.kra")
  next_page_candidate = None
  # print("Current page: "+current_page_dt.strftime(dt_fmt))
  i = 0
  l = list(reversed(sorted(paths)) if prev else sorted(paths))

  for p in l:
    i += 1
    if exact:
      if p == Path(path).joinpath(f"{current_page_dt.strftime(dt_fmt)}.kra"):
        dt = datetime.strptime(Path(p).name, f"{dt_fmt}.kra")
        next_page_candidate = dt
        break
    if p == Path(path).joinpath(f"{current_page_dt.strftime(dt_fmt)}.kra"):
      continue
    try:
      dt = datetime.strptime(Path(p).name, f"{dt_fmt}.kra")
      if prev:
        if current_page_dt > dt:
          next_page_candidate = dt
          break
      else:
        if current_page_dt < dt:
          next_page_candidate = dt
          break
    except Exception as e:
      pass
      # print('sketchbookmode:', str(e))
  if next_page_candidate is not None:
    #Check if next page already open in krita
    # print('cand '+Path(path).joinpath(next_page_candidate.strftime(dt_fmt) + ".kra").as_posix())
    for v in app.activeWindow().views():
      # print('view '+v.document().fileName())
      if v.document().fileName() == Path(path).joinpath(next_page_candidate.strftime(dt_fmt) + ".kra").as_posix():
        doc.setBatchmode(True)
        win.activate()
        win.showView(v)
        v.setVisible()
        currentPreset = view.currentBrushPreset()
        o = f"({len(app.activeWindow().views())}) "
        o += f"{Path(path).stem.upper()} "
        o += generate_id(seed=next_page_candidate.strftime(dt_fmt), separator='-', word_count=3)

        o += f" {len(l) - i + 1 if prev else i}/{len(l)}"
        r = FOLDERS.index(Path(path).name)
        r2 = FOLDERS[r+1 if r+1 < len(FOLDERS) else 0].upper()
        r3 = FOLDERS[r-1 if r-1 >= 0 else -1].upper()
        v.showFloatingMessage(
          f"{r2}\n{o}\n{r3}",  # str(Path(path).joinpath(next_page_candidate.strftime(dt_fmt) + ".kra")),
          QIcon(QPixmap.fromImage(currentPreset.image())),
          3000,
          0
        )
        doc.setBatchmode(False)
        # print("OPEN ALREADY")
        animation_tools_toggle(path.stem)
        history[Path(path).stem.lower()] = next_page_candidate
        return
    doc.setBatchmode(True)
    doc2 = app.openDocument(Path(path).joinpath(f"{next_page_candidate.strftime(dt_fmt)}.kra").as_posix())
    doc.setBatchmode(False)
    animation_tools_toggle(path.stem)
    history[Path(path).stem.lower()] = next_page_candidate
    app.activeWindow().addView(doc2)
  else:
    history[Path(path).stem.lower()] = None
    doc.setBatchmode(True)
    create_new_page2(path)
    doc.setBatchmode(False)
    animation_tools_toggle(path.stem)
    # TODO: automatically close next and previous pages +3 to save memory

class SketchbookMode(Extension):
  def __init__(self, parent):
    self.books = FOLDERS
    self.head_book = "sketches"
    self.head_page = None
    super(SketchbookMode, self).__init__(parent)
    for d in self.books:
      history[d] = None

  def setup(self):
      pass

  def initialize(self):
    pass


  def createActions(self, window):
    def next_page():
      app = Krita.instance()
      win = app.activeWindow()
      current_doc = win.activeView().document()
      if current_doc:
        try:
          dt = datetime.strptime(Path(current_doc.fileName()).name, f"{dt_fmt}.kra")
          get_next_prev_page(Path(current_doc.fileName()).parent, dt)
        except Exception as e:
          pass
          #print(e)
      else:
        get_next_prev_page(Path(current_doc.fileName()).parent, datetime.now()+timedelta(days=9999999), prev=True)

    def prev_page():
      app = Krita.instance()
      win = app.activeWindow()
      current_doc = win.activeView().document()
      if current_doc:
        try:
          dt = datetime.strptime(Path(current_doc.fileName()).name, f"{dt_fmt}.kra")
          get_next_prev_page(Path(current_doc.fileName()).parent, dt, prev=True)
        except Exception as e:
          pass
          #print(e)
      else:
        get_next_prev_page(Path(current_doc.fileName()).parent, datetime.now()+timedelta(days=9999999), prev=True)

    def next_sketchbook():
      print("next_sketchbook")
      app = Krita.instance()
      win = app.activeWindow()
      current_doc = win.activeView().document()
      if current_doc:
        dirname = Path(current_doc.fileName()).parent.stem
        print('dirname', dirname)
        if dirname in self.books:
          i = self.books.index(dirname) + 1
          next_book = self.books[i if i < len(self.books) else i % len(self.books)]
          p = Path(current_doc.fileName()).parent.parent.joinpath(next_book)
          print(p.name.lower())
          if history[p.name.lower()] is None:
            get_next_prev_page(p, datetime.now(), True)
          else:
            get_next_prev_page(p, history[p.name.lower()], True, True)
          # get_last_page(p)
        else:
          p = Path(current_doc.fileName()).parent.parent.joinpath(self.head_book)
          get_next_prev_page(p, datetime.now(), True)
          # get_last_page(p)
        print("next_sketchbook", dirname, p.as_posix())

    def prev_sketchbook():
      print("prev_sketchbook")
      app = Krita.instance()
      win = app.activeWindow()
      current_doc = win.activeView().document()
      if current_doc:
        dirname = Path(current_doc.fileName()).parent.stem
        print('dirname', dirname)
        if dirname in self.books:
          i = self.books.index(dirname) - 1
          next_book = self.books[i if i < len(self.books) else i % len(self.books)]
          p = Path(current_doc.fileName()).parent.parent.joinpath(next_book)
          print(p.name.lower())
          if history[p.name.lower()] is None:
            get_next_prev_page(p, datetime.now(), True)
          else:
            get_next_prev_page(p, history[p.name.lower()], True, True)
          # get_last_page(p)
        else:
          p = Path(current_doc.fileName()).parent.parent.joinpath(self.head_book)
          get_next_prev_page(p, datetime.now(), True)
         # get_last_page(p)
        print("prev_sketchbook", dirname, p.as_posix())

    def toggle_tab_bar():
      from PyQt5.QtWidgets import QStackedWidget, QTabBar, QMdiArea

      qwin = Krita.instance().activeWindow().qwindow()
      for a in qwin.findChildren(QMdiArea):
        e = a.findChild(QTabBar)
        if e:
          e.hide() if e.isVisible() else e.show()
          # print(dir(a.findChild(QTabBar)))
          break

    def fn1():
      view = Krita.instance().activeWindow().activeView()
      canvas = view.canvas()
      canvas.setRotation(canvas.rotation()+90)

    def fn2():
      view = Krita.instance().activeWindow().activeView()
      canvas = view.canvas()
      canvas.mirror()
      canvas.setRotation(180 if canvas.rotation() == 0 else 0)

    action = window.createAction("sketchbook_next_page", str(f"Get next sketchbook page"), "")
    action.triggered.connect(next_page)

    action = window.createAction("sketchbook_prev_page", str(f"Get prev sketchbook page"), "")
    action.triggered.connect(prev_page)

    action = window.createAction("sketchbook_next_sketchbook", str(f"Get next sketchbook type"), "")
    action.triggered.connect(next_sketchbook)

    action = window.createAction("sketchbook_prev_sketchbook", str(f"Get prev sketchbook type"), "")
    action.triggered.connect(prev_sketchbook)

    action = window.createAction("sketchbook_toggle_tab_bar", str(f"Toggle hide tab bar"), "")
    action.triggered.connect(toggle_tab_bar)

    action = window.createAction("rotate_90", str(f"Rotate 90"), "")
    action.triggered.connect(fn1)

    action = window.createAction("flip_vertical", str(f"Flip vert"), "")
    action.triggered.connect(fn2)

    # action = window.createAction("sketchbook_last_page", str(f"Toggle hide tab bar"), "")
    # action.triggered.connect(toggle_tab_bar)