Add support to extract total comment count.

TeamNewPipe · Feb 8, 2022 · 04394ee · 04394ee
1 parent 65129e6
commit 04394ee
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 8 deletions.
diff --git a/extractor/src/main/java/org/schabi/newpipe/extractor/comments/CommentsExtractor.java b/extractor/src/main/java/org/schabi/newpipe/extractor/comments/CommentsExtractor.java
@@ -22,6 +22,13 @@ public boolean isCommentsDisabled() throws ExtractionException {
         return false;
     }
 
+    /**
+     * @return total number of comments.
+     */
+    public int getCommentsCount() throws ExtractionException {
+        return -1;
+    }
+
     @Nonnull
     @Override
     public String getName() throws ParsingException {

diff --git a/extractor/src/main/java/org/schabi/newpipe/extractor/comments/CommentsInfo.java b/extractor/src/main/java/org/schabi/newpipe/extractor/comments/CommentsInfo.java
@@ -48,6 +48,11 @@ public static CommentsInfo getInfo(final CommentsExtractor commentsExtractor)
                 ExtractorHelper.getItemsPageOrLogError(commentsInfo, commentsExtractor);
         commentsInfo.setCommentsDisabled(commentsExtractor.isCommentsDisabled());
         commentsInfo.setRelatedItems(initialCommentsPage.getItems());
+        try {
+            commentsInfo.setCommentsCount(commentsExtractor.getCommentsCount());
+        } catch (Exception e) {
+            commentsInfo.addError(e);
+        }
         commentsInfo.setNextPage(initialCommentsPage.getNextPage());
 
         return commentsInfo;
@@ -72,6 +77,7 @@ public static InfoItemsPage<CommentsInfoItem> getMoreItems(
 
     private transient CommentsExtractor commentsExtractor;
     private boolean commentsDisabled = false;
+    private int commentsCount;
 
     public CommentsExtractor getCommentsExtractor() {
         return commentsExtractor;
@@ -82,19 +88,37 @@ public void setCommentsExtractor(final CommentsExtractor commentsExtractor) {
     }
 
     /**
-     * @apiNote Warning: This method is experimental and may get removed in a future release.
      * @return <code>true</code> if the comments are disabled otherwise <code>false</code> (default)
+     * @apiNote Warning: This method is experimental and may get removed in a future release.
      * @see CommentsExtractor#isCommentsDisabled()
      */
     public boolean isCommentsDisabled() {
         return commentsDisabled;
     }
 
     /**
-     * @apiNote Warning: This method is experimental and may get removed in a future release.
      * @param commentsDisabled <code>true</code> if the comments are disabled otherwise <code>false</code>
+     * @apiNote Warning: This method is experimental and may get removed in a future release.
      */
     public void setCommentsDisabled(final boolean commentsDisabled) {
         this.commentsDisabled = commentsDisabled;
     }
+
+    /**
+     * Returns the total number of comments.
+     *
+     * @return totalComments
+     */
+    public int getCommentsCount() {
+        return commentsCount;
+    }
+
+    /**
+     * Sets the total number of comments.
+     *
+     * @param commentsCount
+     */
+    public void setCommentsCount(int commentsCount) {
+        this.commentsCount = commentsCount;
+    }
 }
diff --git a/...va/org/schabi/newpipe/extractor/services/youtube/extractors/YoutubeCommentsExtractor.java b/...va/org/schabi/newpipe/extractor/services/youtube/extractors/YoutubeCommentsExtractor.java
@@ -1,7 +1,6 @@
 package org.schabi.newpipe.extractor.services.youtube.extractors;
 
-import static org.schabi.newpipe.extractor.services.youtube.YoutubeParsingHelper.getJsonPostResponse;
-import static org.schabi.newpipe.extractor.services.youtube.YoutubeParsingHelper.prepareDesktopJsonBuilder;
+import static org.schabi.newpipe.extractor.services.youtube.YoutubeParsingHelper.*;
 import static org.schabi.newpipe.extractor.utils.Utils.UTF_8;
 import static org.schabi.newpipe.extractor.utils.Utils.isNullOrEmpty;
 
@@ -29,6 +28,7 @@
 import com.grack.nanojson.JsonArray;
 import com.grack.nanojson.JsonObject;
 import com.grack.nanojson.JsonWriter;
+import org.schabi.newpipe.extractor.utils.Utils;
 
 public class YoutubeCommentsExtractor extends CommentsExtractor {
 
@@ -45,6 +45,7 @@ public class YoutubeCommentsExtractor extends CommentsExtractor {
      */
     @SuppressWarnings("OptionalUsedAsFieldOrParameterType")
     private Optional<Boolean> optCommentsDisabled = Optional.empty();
+    private JsonObject ajaxJson;
 
     public YoutubeCommentsExtractor(
             final StreamingService service,
@@ -175,16 +176,15 @@ public InfoItemsPage<CommentsInfoItem> getPage(final Page page)
                 .done())
                 .getBytes(UTF_8);
 
-        final JsonObject ajaxJson = getJsonPostResponse("next", body, localization);
+        this.ajaxJson = getJsonPostResponse("next", body, localization);
 
         final CommentsInfoItemsCollector collector = new CommentsInfoItemsCollector(
                 getServiceId());
-        collectCommentsFrom(collector, ajaxJson);
+        collectCommentsFrom(collector);
         return new InfoItemsPage<>(collector, getNextPage(ajaxJson));
     }
 
-    private void collectCommentsFrom(final CommentsInfoItemsCollector collector,
-                                     @Nonnull final JsonObject ajaxJson) throws ParsingException {
+    private void collectCommentsFrom(final CommentsInfoItemsCollector collector) throws ParsingException {
 
         final JsonArray onResponseReceivedEndpoints = ajaxJson.getArray(
                 "onResponseReceivedEndpoints");
@@ -257,4 +257,17 @@ public boolean isCommentsDisabled() throws ExtractionException {
 
         return optCommentsDisabled.get();
     }
+
+    @Override
+    public int getCommentsCount() throws ExtractionException {
+        final JsonObject commentsHeaderRenderer = ajaxJson
+                .getArray("onResponseReceivedEndpoints").getObject(0)
+                .getObject("reloadContinuationItemsCommand")
+                .getArray("continuationItems").getObject(0)
+                .getObject("commentsHeaderRenderer");
+
+        final String text = getTextFromObject(commentsHeaderRenderer.getObject("countText"));
+
+        return Integer.parseInt(Utils.removeNonDigitCharacters(text));
+    }
 }
diff --git a/...test/java/org/schabi/newpipe/extractor/services/youtube/YoutubeCommentsExtractorTest.java b/...test/java/org/schabi/newpipe/extractor/services/youtube/YoutubeCommentsExtractorTest.java
@@ -87,6 +87,7 @@ private boolean getCommentsFromCommentsInfoHelper(String url) throws IOException
         @Test
         public void testGetCommentsAllData() throws IOException, ExtractionException {
             InfoItemsPage<CommentsInfoItem> comments = extractor.getInitialPage();
+            assertTrue(extractor.getCommentsCount() > 5); // at least 5 comments
 
             DefaultTests.defaultTestListOfItems(YouTube, comments.getItems(), comments.getErrors());
             for (CommentsInfoItem c : comments.getItems()) {
@@ -336,5 +337,11 @@ public void testGetCommentsFirstReplies() throws IOException, ExtractionExceptio
             assertEquals("First", replies.getItems().get(0).getCommentText(),
                     "First reply comment did not match");
         }
+
+        @Test
+        public void testCommentsCount() throws IOException, ExtractionException {
+            extractor.getInitialPage(); // Needs to be called first
+            assertTrue(extractor.getCommentsCount() > 18800);
+        }
     }
 }