Akka Stream operator to split a stream by type

MessageByType.scala

import akka.NotUsed
import akka.http.scaladsl.model.ws
import akka.http.scaladsl.model.ws.{BinaryMessage, Message, TextMessage}
import akka.stream.scaladsl.{Broadcast, Flow, GraphDSL}
import akka.stream.{FanOutShape, Graph, Outlet}

import scala.collection.immutable

/**
  * Partition websocket messages into binary and text messages.
  *
  * The flow takes websocket message and partitions them into two outlets, one
  * for binary messages, the other for text messages.
  */
object MessageByType {

  def apply(): Graph[Shape, NotUsed] =
    GraphDSL.create() { implicit b =>
      import GraphDSL.Implicits._

      val bcast = b.add(Broadcast[ws.Message](2))
      val collectText = b.add(Flow[ws.Message].collectType[ws.TextMessage])
      val collectBinary = b.add(Flow[ws.Message].collectType[ws.BinaryMessage])

      bcast.out(0) ~> collectBinary.in
      bcast.out(1) ~> collectText.in

      new Shape(
        FanOutShape
          .Ports(
            bcast.in,
            immutable.Seq(collectText.out, collectBinary.out)
          )
      )
    }

  /**
    * Shape of the partition by type graph.
    */
  class Shape(init: FanOutShape.Init[ws.Message] = FanOutShape.Name("MessageByType"))
      extends FanOutShape[ws.Message](init) {
    override protected def construct(init: FanOutShape.Init[Message]): FanOutShape[Message] =
      new Shape(init)

    val textOut: Outlet[TextMessage] = newOutlet[ws.TextMessage]("textOut")
    val binaryOut: Outlet[BinaryMessage] = newOutlet[ws.BinaryMessage]("binaryOut")
  }
}