<commit_before>package com.karasiq.shadowcloud.streams.metadata

import akka.NotUsed
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.stream.scaladsl.{Flow, Source}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import akka.util.ByteString

import com.karasiq.shadowcloud.metadata.MimeDetector

object MimeDetectorStream {
  def apply(detector: MimeDetector, fileName: String, probeSize: Int): Flow[ByteString, String, NotUsed] = {
    Flow.fromGraph(new MimeDetectorStream(detector, fileName, probeSize))
      .recoverWithRetries(1, { case _ ⇒ Source.single(MimeDetector.DefaultMime) })
      .orElse(Source.single(MimeDetector.DefaultMime))
      .take(1)
  }
}

private final class MimeDetectorStream(detector: MimeDetector, fileName: String, probeSize: Int)
  extends GraphStage[FlowShape[ByteString, String]] {

  val inlet = Inlet[ByteString]("MimeDetectorStream.in")
  val outlet = Outlet[String]("MimeDetectorStream.out")
  val shape = FlowShape(inlet, outlet)

  def createLogic(inheritedAttributes: Attributes) = new GraphStageLogic(shape) with InHandler with OutHandler {
    private[this] var buffer = ByteString.empty

    private[this] def pushContentType(): Unit = {
      if (buffer.length >= probeSize || isClosed(inlet)) {
        val contentType = detector.getMimeType(fileName, buffer.take(probeSize))
        buffer = ByteString.empty
        push(outlet, contentType.getOrElse(MimeDetector.DefaultMime))
        completeStage()
      } else {
        tryPull(inlet)
      }
    }

    def onPull(): Unit = {
      tryPull(inlet)
    }

    def onPush(): Unit = {
      val element = grab(inlet)
      buffer ++= element
      pushContentType()
    }

    setHandlers(inlet, outlet, this)
  }
}<commit_msg>Mime fixes<commit_after>package com.karasiq.shadowcloud.streams.metadata

import akka.NotUsed
import akka.stream.{Attributes, FlowShape, Inlet, Outlet}
import akka.stream.scaladsl.{Flow, Source}
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler, OutHandler}
import akka.util.ByteString

import com.karasiq.shadowcloud.metadata.MimeDetector

object MimeDetectorStream {
  def apply(detector: MimeDetector, fileName: String, probeSize: Int): Flow[ByteString, String, NotUsed] = {
    val defaultMimeSource = Source.single(MimeDetector.DefaultMime)
    Flow.fromGraph(new MimeDetectorStream(detector, fileName, probeSize))
      .recoverWithRetries(1, { case _ ⇒ defaultMimeSource })
      .orElse(defaultMimeSource)
  }
}

private final class MimeDetectorStream(detector: MimeDetector, fileName: String, probeSize: Int)
  extends GraphStage[FlowShape[ByteString, String]] {

  val inBytes = Inlet[ByteString]("MimeDetectorStream.in")
  val outContentType = Outlet[String]("MimeDetectorStream.out")
  val shape = FlowShape(inBytes, outContentType)

  def createLogic(inheritedAttributes: Attributes) = new GraphStageLogic(shape) with InHandler with OutHandler {
    private[this] var buffer = ByteString.empty

    private[this] def pushContentType(): Unit = {
      if (buffer.length >= probeSize || isClosed(inBytes)) {
        val contentType = detector.getMimeType(fileName, buffer/*.take(probeSize)*/)
        buffer = ByteString.empty
        contentType.foreach(push(outContentType, _))
        completeStage()
      } else {
        tryPull(inBytes)
      }
    }

    def onPull(): Unit = {
      tryPull(inBytes)
    }

    def onPush(): Unit = {
      val element = grab(inBytes)
      buffer ++= element
      pushContentType()
    }

    override def onUpstreamFinish(): Unit = {
      pushContentType()
    }

    setHandlers(inBytes, outContentType, this)
  }
}
{'fraction_non_alphanumeric': 0.10714285714285714, 'fraction_numerical': 0.0008178844056706652, 'mean_word_length': 4.124301675977653, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2773875', 'n_tokens_mistral': 1175, 'n_tokens_neox': 1154, 'n_words': 233}