カスタムセルの高さを動的に変更したい

Question

前提・実現したいこと

カスタムセルに画像を0〜2つ表示させ、その個数にセルの高さを対応させたいと考えています。

条件
・カスタムセルに0〜2つの画像が表示される。
・画像の大きさは、適当に(width: 100, height: 100)としている。
・画像が複数ある場合は、縦に並べる。

現段階での解決方法は、
①カスタムセルに画像をセットした後、その画像の高さの合計をカスタムセル内のcellHeightプロパティに保存。
②cellForRowAt内でカスタムセルからcellHeightの値を取得。
③tableView.rowHeightにcellHeightの値を設定。
としております。

理想のコードとしては、
カスタムセルから画像がはみ出さずに、カスタムセルの高さが自動的に変更される。
というような、コードを書きたいと思っております。

現時点のコードのような書き方以外で、セルごとに高さを変えられるような書き方はありますでしょうか？
よろしくお願い致します(＞人＜;)

該当のソースコード

class CustomTableViewCell: UITableViewCell {
    var firstImageView: UIImageView!
    var secondImageView: UIImageView!
    var cellHeight: CGFloat = 0

    override init(style: UITableViewCell.CellStyle, reuseIdentifier: String?) {
        super.init(style: style, reuseIdentifier: reuseIdentifier)
        firstImageView = UIImageView()
        secondImageView = UIImageView()
        
        firstImageView.translatesAutoresizingMaskIntoConstraints = false
        secondImageView.translatesAutoresizingMaskIntoConstraints = false

        firstImageView.contentMode = .scaleAspectFit
        secondImageView.contentMode = .scaleAspectFit
        
        let userImageSize = CGSize(width: 100, height: 100)
        
        firstImageView.bounds = CGRect(x: 0, y: 0, width: userImageSize.width, height: userImageSize.height)
        secondImageView.bounds = firstImageView.bounds

        self.addSubview(firstImageView)
        self.addSubview(secondImageView)
    }
    
    override func layoutSubviews() {
        guard firstImageView.image != nil else {
            return
        }
        firstImageView.widthAnchor.constraint(equalToConstant: 100).isActive = true
        firstImageView.heightAnchor.constraint(equalToConstant: 100).isActive = true
        firstImageView.leadingAnchor.constraint(equalTo: self.leadingAnchor).isActive = true
        firstImageView.topAnchor.constraint(equalTo: self.topAnchor).isActive = true
        cellHeight += firstImageView.frame.height
        
        guard secondImageView.image != nil else {
            return
        }
        secondImageView.widthAnchor.constraint(equalToConstant: 100).isActive = true
        secondImageView.heightAnchor.constraint(equalToConstant: 100).isActive = true
        secondImageView.leadingAnchor.constraint(equalTo: self.leadingAnchor).isActive = true
        secondImageView.topAnchor.constraint(equalTo: firstImageView.bottomAnchor).isActive = true
        cellHeight += secondImageView.frame.height
    }
}

    func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
        let cell = tableView.dequeueReusableCell(withIdentifier: "cell") as! CustomTableViewCell
        let item = dataArray[indexPath.row]
        cell.setCell(first: item.firstImage, second: item.secondImage)
        cell.layoutIfNeeded()
        tableView.rowHeight = cell.cellHeight
        return cell
    }

試したこと

UITableView.automaticDimenssionを使ってみたのですが上手くいきませんでした。

Answer 1

tableViewのheightForRowAtメソッドで出来ますね。

override func tableView(_ tableView: UITableView, heightForRowAt indexPath: IndexPath) -> CGFloat {
       if dataArray.count >= 2 {
          return 100 //画像の数が2つ以上だったら100を指定
       } else {
          return 50 //それ以外だったら50を指定
       }    
    }