From f662194e4dcdf29e1317b832d90c1fdf9a708dd8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=8C=E8=B4=A7=E7=88=B1=E5=90=83=E7=99=BD=E8=90=9D?=
 =?UTF-8?q?=E5=8D=9C?= <smith3816@gmail.com>
Date: Wed, 7 Aug 2024 16:51:35 +0800
Subject: [PATCH] docs: fix limit demo (#50287)

---
 .../date-picker/demo/select-in-range.tsx      | 40 ++++++++++++++++---
 1 file changed, 34 insertions(+), 6 deletions(-)

diff --git a/components/date-picker/demo/select-in-range.tsx b/components/date-picker/demo/select-in-range.tsx
index f503d39b51..3d02a8457d 100644
--- a/components/date-picker/demo/select-in-range.tsx
+++ b/components/date-picker/demo/select-in-range.tsx
@@ -1,24 +1,52 @@
 import React from 'react';
 import { DatePicker, Space, Typography } from 'antd';
 import type { DatePickerProps } from 'antd';
+import type { Dayjs } from 'dayjs';
 
 const { RangePicker } = DatePicker;
 
+const getYearMonth = (date: Dayjs) => date.year() * 12 + date.month();
+
 // Disabled 7 days from the selected date
-const disabled7DaysDate: DatePickerProps['disabledDate'] = (current, { from }) => {
+const disabled7DaysDate: DatePickerProps['disabledDate'] = (current, { from, type }) => {
   if (from) {
-    return Math.abs(current.diff(from, 'days')) >= 7;
+    const minDate = from.add(-6, 'days');
+    const maxDate = from.add(6, 'days');
+
+    switch (type) {
+      case 'year':
+        return current.year() < minDate.year() || current.year() > maxDate.year();
+
+      case 'month':
+        return (
+          getYearMonth(current) < getYearMonth(minDate) ||
+          getYearMonth(current) > getYearMonth(maxDate)
+        );
+
+      default:
+        return Math.abs(current.diff(from, 'days')) >= 7;
+    }
   }
 
   return false;
 };
 
 // Disabled 6 months from the selected date
-const disabled6MonthsDate: DatePickerProps['disabledDate'] = (current, { from }) => {
+const disabled6MonthsDate: DatePickerProps['disabledDate'] = (current, { from, type }) => {
   if (from) {
-    const curMonth = current.year() * 12 + current.month();
-    const fromMonth = from.year() * 12 + from.month();
-    return Math.abs(fromMonth - curMonth) >= 6;
+    const minDate = from.add(-5, 'months');
+    const maxDate = from.add(5, 'months');
+
+    switch (type) {
+      case 'year':
+        return current.year() < minDate.year() || current.year() > maxDate.year();
+
+      default:
+        return (
+          getYearMonth(current) < getYearMonth(minDate) ||
+          getYearMonth(current) > getYearMonth(maxDate)
+        );
+    }
   }
 
   return false;